2. Extract features.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73e9ab10",
   "metadata": {},
   "outputs": [],
   "source": [
    "#default_exp extract_features"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "baa2acee",
   "metadata": {},
   "source": [
    "# Extract the image features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3a3852f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "from fastcore.script import *\n",
    "from fastai.vision.all import *\n",
    "from fastai.callback.wandb import *\n",
    "from self_supervised.augmentations import *\n",
    "from self_supervised.layers import *\n",
    "from self_supervised.vision.barlow_twins import *\n",
    "\n",
    "# fixup depreciation\n",
    "from kornia import augmentation as korniatfm\n",
    "korniatfm.GaussianBlur = korniatfm.RandomGaussianBlur"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d746195b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@call_parse\n",
    "def extract_features(\n",
    "        path:Path, # path\n",
    "        backbone:str    = 'resnet18', # backbone\n",
    "        pretrained:bool = True, # pretrained\n",
    "        size:int        = 224, # size\n",
    "        projection_size = 2048, # projection_size\n",
    "        lr:float        = 5e-4, # lr\n",
    "        epochs:int      = 5, # epochs\n",
    "        model_file:str  = None, # model .pth file\n",
    "    ):\n",
    "    \"Pretrain a model on images in `path`.\"\n",
    "    if model_file is None:\n",
    "        model_path = path/'models'\n",
    "        models = list(model_path.glob('*.pth'))\n",
    "        if len(models) == 0:\n",
    "            raise Exception(f\"No models found in: {model_path}!\")\n",
    "        elif len(models) > 1:\n",
    "            msg = f\"Found multiple models in: {model_path}, please pass one of:\"\n",
    "            for model in models:\n",
    "                msg += f\"\\n  --model_file {model}\"\n",
    "            raise Exception(msg)\n",
    "        model_file = models[0]\n",
    "\n",
    "    fastai_encoder = create_encoder('resnet18', pretrained=pretrained)\n",
    "    model = create_barlow_twins_model(fastai_encoder, projection_size=projection_size, hidden_size=projection_size)\n",
    "\n",
    "    dls = ImageDataLoaders.from_name_func(path, get_image_files(path), lambda x: False, item_tfms=Resize(size), valid_pct=0)\n",
    "    fnames = get_image_files(path)\n",
    "    test_dl = dls.test_dl(fnames)\n",
    "\n",
    "    load_model(model_file, model, None, with_opt=False)\n",
    "    if torch.cuda.is_available():\n",
    "        model = model.cuda()\n",
    "    \n",
    "    features = []\n",
    "    with torch.no_grad():\n",
    "        for (dx,) in progress_bar(test_dl):\n",
    "            features.append(model(dx).detach())\n",
    "    raw_features = torch.cat(features)\n",
    "    features = raw_features / torch.linalg.vector_norm(raw_features, dim=1, keepdims=True)\n",
    "\n",
    "    with open(path/(Path(model_file).with_suffix('').name+'-features.pkl'), 'wb') as f:\n",
    "        pickle.dump(dict(fnames = list(fnames), features = features), f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d26419f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "<style>\n",
       "    /* Turns off some styling */\n",
       "    progress {\n",
       "        /* gets rid of default border in Firefox and Opera. */\n",
       "        border: none;\n",
       "        /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "        background-size: auto;\n",
       "    }\n",
       "    .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "        background: #F44336;\n",
       "    }\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      <progress value='142' class='' max='142' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      100.00% [142/142 04:37<00:00]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "extract_features(Path('../women-in-tshirts-256/'), model_file='../women-in-tshirts-256/models/barlow-twins-resnet18-pretrained-224-5e-proj2048-lr0.5e-3.pth')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}