run_inference_on_example_images.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-05-25T21:04:01.881647Z",
     "start_time": "2024-05-25T21:04:01.864152Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import torch.utils.data\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import utils\n",
    "import vision_transformer as vits\n",
    "from constants import Column\n",
    "import augmentations as aug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [
    "# Parameters\n",
    "arch = 'vit_small'\n",
    "patch_size = 16\n",
    "in_channels = 4\n",
    "n_last_blocks = 4\n",
    "avgpool_patchtokens = False\n",
    "pretrained_weights = './checkpoints/pretrained_cross_batch_n16_guide0_ntc_norm.ckpt'\n",
    "checkpoint_key = 'teacher'"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-05-25T20:44:01.222234Z",
     "start_time": "2024-05-25T20:44:01.187117Z"
    }
   },
   "id": "d9866c91f773c5fb"
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pretrained weights found at ./checkpoints/pretrained_cross_batch_n16_guide0_ntc_norm.ckpt and loaded with msg: <All keys matched successfully>\n"
     ]
    }
   ],
   "source": [
    "model = vits.__dict__[arch](patch_size=patch_size, num_classes=0, in_chans=in_channels)\n",
    "embed_dim = model.embed_dim * (n_last_blocks + int(avgpool_patchtokens))\n",
    "\n",
    "# Run inference\n",
    "if torch.cuda.is_available():\n",
    "    model.cuda()\n",
    "model.eval()\n",
    "# # load weights to evaluate\n",
    "utils.load_pretrained_weights(model, pretrained_weights, checkpoint_key, arch, patch_size)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-05-25T20:44:06.474775Z",
     "start_time": "2024-05-25T20:44:02.054700Z"
    }
   },
   "id": "4bb0bee742bea2b1"
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Run inference on example images"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "843daebbb6d7ea29"
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "====== zscore =====\n"
     ]
    }
   ],
   "source": [
    "files = os.listdir('./example_data/images')\n",
    "# here we used zscore as an example\n",
    "# we suggest to use the statistics calcualated from NTCs in practice\n",
    "normalizer = aug.Normalization(method='zscore')\n",
    "image_list = []\n",
    "keys =[]\n",
    "for file in files:\n",
    "    keys.append(file.split('.npy')[0])\n",
    "    image = np.load(os.path.join('./example_data/images', file))\n",
    "    image_norm = normalizer(image)\n",
    "    image_list.append(image_norm)\n",
    "images = torch.Tensor(np.stack(image_list, axis=0))"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-05-25T21:06:35.318230Z",
     "start_time": "2024-05-25T21:06:35.297575Z"
    }
   },
   "id": "fdbe4e4765ef798a"
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    # image, key\n",
    "    if torch.cuda.is_available():\n",
    "        input = images.cuda(non_blocking=True)\n",
    "    else:\n",
    "        input = images\n",
    "\n",
    "    if \"vit\" in arch:\n",
    "        intermediate_output = model.get_intermediate_layers(input, n_last_blocks)\n",
    "        output = torch.cat([x[:, 0] for x in intermediate_output], dim=-1)\n",
    "        if avgpool_patchtokens:\n",
    "            output = torch.cat((output.unsqueeze(-1),\n",
    "                                torch.mean(intermediate_output[-1][:, 1:], dim=1).unsqueeze(-1)), dim=-1)\n",
    "            output = output.reshape(output.shape[0], -1)\n",
    "    else:\n",
    "        output = model(input)\n",
    "\n",
    "    if torch.cuda.is_available():\n",
    "        output = output.cpu().numpy()\n",
    "    else:\n",
    "        output = output.numpy()\n",
    "\n",
    "embed_dim = output.shape[-1]\n",
    "output_df = pd.DataFrame(output, columns=[f'feature_{idx}' for idx in range(embed_dim)])\n",
    "\n",
    "# Add metadata\n",
    "output_df['key'] = keys\n",
    "output_df[[Column.plate.value, Column.well.value, Column.tile.value, Column.gene.value,\n",
    "           Column.sgRNA.value, 'meta_df_index']] = \\\n",
    "    output_df.apply(lambda x: pd.Series(str(x['key']).split(';')), axis=1)\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-05-25T21:07:12.494029Z",
     "start_time": "2024-05-25T21:07:12.475287Z"
    }
   },
   "id": "44f2b824d816322b"
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Run inference on example LMDB dataset"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "36e0e4df90e650a"
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "outputs": [
    {
     "data": {
      "text/plain": "<module 'dataset' from '/home/yaoh11/cellpaint/set-dino/dataset.py'>"
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import dataset\n",
    "\n",
    "import importlib\n",
    "importlib.reload(dataset)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-05-25T21:14:38.879369Z",
     "start_time": "2024-05-25T21:14:38.854478Z"
    }
   },
   "id": "4009c577af8fe229"
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "====== zscore =====\n"
     ]
    }
   ],
   "source": [
    "dataset_path = './example_data/lmdb_dataset'\n",
    "crop_size = 96\n",
    "# Load metadata\n",
    "df = pd.read_csv('./example_data/metadata.csv', index_col=0)\n",
    "\n",
    "# Create the dataset and data loader\n",
    "normalizer = aug.Normalization(method='zscore')\n",
    "ds = dataset.InferenceDataset(df, dataset_path=dataset_path, crop_size=crop_size, normalizer=normalizer)\n",
    "\n",
    "dataloader = torch.utils.data.DataLoader(\n",
    "    dataset=ds,\n",
    "    batch_size=8,\n",
    "    shuffle=False,\n",
    "    num_workers=0,\n",
    "    drop_last=False,\n",
    "    )"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-05-25T21:14:41.803257Z",
     "start_time": "2024-05-25T21:14:41.774280Z"
    }
   },
   "id": "f885ccba739050b"
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "outputs": [],
   "source": [
    "output_list = []\n",
    "key_list = []\n",
    "with torch.no_grad():\n",
    "    for idx, (images, keys) in enumerate(dataloader):\n",
    "        if torch.cuda.is_available():\n",
    "            input = images.cuda(non_blocking=True)\n",
    "        else:\n",
    "            input = images\n",
    "\n",
    "        if \"vit\" in arch:\n",
    "            intermediate_output = model.get_intermediate_layers(input, n_last_blocks)\n",
    "            output = torch.cat([x[:, 0] for x in intermediate_output], dim=-1)\n",
    "            if avgpool_patchtokens:\n",
    "                output = torch.cat((output.unsqueeze(-1),\n",
    "                                    torch.mean(intermediate_output[-1][:, 1:], dim=1).unsqueeze(-1)), dim=-1)\n",
    "                output = output.reshape(output.shape[0], -1)\n",
    "        else:\n",
    "            output = model(input)\n",
    "\n",
    "        if torch.cuda.is_available():\n",
    "            output_list.append(output.cpu().numpy())\n",
    "        else:\n",
    "            output_list.append(output.numpy())\n",
    "        key_list.extend(list(keys))\n",
    "\n",
    "output_all = np.concatenate(output_list, axis=0)\n",
    "embed_dim = output.shape[-1]\n",
    "output_df = pd.DataFrame(output_all, columns=[f'feature_{idx}' for idx in range(embed_dim)])\n",
    "\n",
    "# Add metadata\n",
    "output_df['key'] = key_list\n",
    "output_df[[Column.plate.value, Column.well.value, Column.tile.value, Column.gene.value,\n",
    "           Column.sgRNA.value, 'meta_df_index']] = \\\n",
    "    output_df.apply(lambda x: pd.Series(str(x['key']).split(';')), axis=1)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-05-25T21:14:44.606860Z",
     "start_time": "2024-05-25T21:14:44.585824Z"
    }
   },
   "id": "5114c0521f506a35"
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "set-dino",
   "language": "python",
   "display_name": "set-dino"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}