modality_gap_cloome.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exploring CLOOME with Amumo \n",
    "Humer et al. 2023\n",
    "\n",
    "Github Repo: https://github.com/ginihumer/Amumo\n",
    "\n",
    "Interactive Article: https://jku-vds-lab.at/amumo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install git+https://github.com/ginihumer/Amumo.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import sys\n",
    "import os\n",
    "import torch\n",
    "sys.path.insert(0, os.path.abspath(\"src/\"))\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Christina\\AppData\\Roaming\\Python\\Python39\\site-packages\\umap\\distances.py:1063: NumbaDeprecationWarning: \u001b[1mThe 'nopython' keyword argument was not supplied to the 'numba.jit' decorator. The implicit default value for this argument is currently False, but it will be changed to True in Numba 0.59.0. See https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit for details.\u001b[0m\n",
      "  @numba.jit()\n",
      "C:\\Users\\Christina\\AppData\\Roaming\\Python\\Python39\\site-packages\\umap\\distances.py:1071: NumbaDeprecationWarning: \u001b[1mThe 'nopython' keyword argument was not supplied to the 'numba.jit' decorator. The implicit default value for this argument is currently False, but it will be changed to True in Numba 0.59.0. See https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit for details.\u001b[0m\n",
      "  @numba.jit()\n",
      "C:\\Users\\Christina\\AppData\\Roaming\\Python\\Python39\\site-packages\\umap\\distances.py:1086: NumbaDeprecationWarning: \u001b[1mThe 'nopython' keyword argument was not supplied to the 'numba.jit' decorator. The implicit default value for this argument is currently False, but it will be changed to True in Numba 0.59.0. See https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit for details.\u001b[0m\n",
      "  @numba.jit()\n",
      "C:\\Users\\Christina\\AppData\\Roaming\\Python\\Python39\\site-packages\\umap\\umap_.py:660: NumbaDeprecationWarning: \u001b[1mThe 'nopython' keyword argument was not supplied to the 'numba.jit' decorator. The implicit default value for this argument is currently False, but it will be changed to True in Numba 0.59.0. See https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit for details.\u001b[0m\n",
      "  @numba.jit()\n"
     ]
    }
   ],
   "source": [
    "import amumo\n",
    "from amumo import model as am_model\n",
    "from amumo import data as am_data\n",
    "from amumo import widgets as am_widgets\n",
    "from amumo import utils as am_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "basepath = './'\n",
    "datapath = os.path.join(basepath, \"data\")\n",
    "batch_size = 100\n",
    "seed = 31415"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "\n",
    "# Data Helpers\n",
    "def get_data_helper(dataset, filters=[], method=any):\n",
    "    all_images, all_prompts = dataset.get_filtered_data(filters, method=method)\n",
    "    print(len(all_images))\n",
    "\n",
    "    dataset_name = dataset.name\n",
    "    if len(filters) > 0:\n",
    "        dataset_name = dataset_name + '_filter-' + method.__name__ + '_' + '-'.join(filters)\n",
    "    else:\n",
    "        dataset_name = dataset_name + '_size-%i'%len(all_images)\n",
    "\n",
    "    return all_images, all_prompts, dataset_name\n",
    "\n",
    "\n",
    "def reshape_image(arr):\n",
    "    c, h, w = arr.shape\n",
    "    reshaped_image = np.empty((h, w, c))\n",
    "\n",
    "    reshaped_image[:,:,0] = arr[0]\n",
    "    reshaped_image[:,:,1] = arr[1]\n",
    "    reshaped_image[:,:,2] = arr[2]\n",
    "\n",
    "    reshaped_pil = Image.fromarray(reshaped_image.astype(\"uint8\"))\n",
    "\n",
    "    return reshaped_pil\n",
    "\n",
    "\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import rdFMCS\n",
    "import io\n",
    "class MoleculeType(am_data.TextType):\n",
    "    name = \"Molecule\"\n",
    "\n",
    "    def __init__(self, data) -> None:\n",
    "        # data is a list of SMILES\n",
    "        super().__init__(data)\n",
    "\n",
    "    def getMinSummary(self, ids):\n",
    "        # retrieve MCS of mols\n",
    "        if len(ids) == 1:\n",
    "            return self.data[ids[0]]\n",
    "\n",
    "        mols = [Chem.MolFromSmiles(smiles) for smiles in self.data[ids]]\n",
    "        mcs = rdFMCS.FindMCS(mols)\n",
    "        mcs_smiles = Chem.MolToSmiles(Chem.MolFromSmarts(mcs.smartsString))\n",
    "        return mcs_smiles\n",
    "    \n",
    "    def getVisItem(self, idx):\n",
    "        output_img = io.BytesIO()\n",
    "        img = Chem.Draw.MolToImage(Chem.MolFromSmiles(self.data[idx]))\n",
    "        img.resize((300,300)).save(output_img, format='JPEG')\n",
    "        return output_img\n",
    "    \n",
    "\n",
    "class BioImageType(am_data.ImageType):\n",
    "    name = \"Bio Image\"\n",
    "\n",
    "    def __init__(self, data) -> None:\n",
    "        super().__init__(data)\n",
    "\n",
    "\n",
    "class CLOOMDataset_Dataset(am_data.DatasetInterface):\n",
    "    name='CLOOMDataset'\n",
    "\n",
    "    def __init__(self, path, seed=31415, batch_size = 100):\n",
    "        super().__init__(path, seed, batch_size)\n",
    "\n",
    "        self.MODE1_Type = BioImageType\n",
    "        self.MODE2_Type = MoleculeType\n",
    "\n",
    "        mol_index_file = os.path.join(path, \"cellpainting-unique-molecule.csv\")\n",
    "        img_index_file = os.path.join(path, \"cellpainting-all-imgpermol.csv\")\n",
    "        images_arr = os.path.join(path, \"subset_npzs_dict_.npz\")\n",
    "\n",
    "\n",
    "        # molecule smiles\n",
    "        all_molecules = pd.read_csv(mol_index_file)\n",
    "        all_molecules.rename(columns={\"SAMPLE_KEY\": \"SAMPLE_KEY_mol\"}, inplace=True)\n",
    "        # microscopy images\n",
    "        all_microscopies = pd.read_csv(img_index_file)\n",
    "        all_microscopies.rename(columns={\"SAMPLE_KEY\": \"SAMPLE_KEY_img\"}, inplace=True)\n",
    "        # join the two dataframes\n",
    "        cloome_data = pd.merge(all_molecules[[\"SAMPLE_KEY_mol\", \"SMILES\"]], all_microscopies[[\"SAMPLE_KEY_img\", \"SMILES\"]], on=\"SMILES\", how=\"inner\")\n",
    "        \n",
    "        # subsample data\n",
    "        self.subset_idcs = self._get_random_subsample(len(cloome_data))\n",
    "        self.dataset = cloome_data.iloc[self.subset_idcs]\n",
    "        \n",
    "        self.all_prompts = self.dataset[\"SMILES\"].values\n",
    "\n",
    "        # microscopy images TODO... load images on demand with a custom image loader\n",
    "        all_microscopies = pd.read_csv(img_index_file)\n",
    "        images_dict = np.load(images_arr, allow_pickle = True)\n",
    "        all_images = []\n",
    "        for img_id in self.dataset[\"SAMPLE_KEY_img\"]:\n",
    "            image = images_dict[f\"{img_id}.npz\"]\n",
    "            im = reshape_image(image)\n",
    "            all_images.append(im)\n",
    "\n",
    "        self.all_images = np.array(all_images)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Christina\\AppData\\Local\\Temp\\ipykernel_28556\\2905515878.py:102: FutureWarning: The input object of type 'Image' is an array-like implementing one of the corresponding protocols (`__array__`, `__array_interface__` or `__array_struct__`); but not a sequence (or 0-D). In the future, this object will be coerced as if it was first converted using `np.array(obj)`. To retain the old behaviour, you have to either modify the type 'Image', or assign to an empty array created with `np.empty(correct_shape, dtype=object)`.\n",
      "  self.all_images = np.array(all_images)\n",
      "C:\\Users\\Christina\\AppData\\Local\\Temp\\ipykernel_28556\\2905515878.py:102: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.\n",
      "  self.all_images = np.array(all_images)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'CLOOMDataset_size-100'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load Data\n",
    "dataset_cloome = CLOOMDataset_Dataset(datapath,seed,batch_size)\n",
    "cloome_images, cloome_molecules, cloome_dataset_name = get_data_helper(dataset_cloome, filters=[], method=any)\n",
    "cloome_dataset_name\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<_io.BytesIO at 0x20775295630>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cloome_molecules.getVisItem(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PrecalculatedModel(am_model.CLIPModelInterface):\n",
    "    model_name = 'precalculated'\n",
    "\n",
    "    def __init__(self, name, dataset_name, modality1_features, modality2_features, logit_scale=torch.tensor(0)) -> None:\n",
    "        # this class is a workaround for precalculated features\n",
    "        # it just saves the features as cached files so that the \"encode_image\" and \"encode_text\" methods are not called\n",
    "        self.available_models = [name]\n",
    "        super().__init__(name, device='cpu')\n",
    "        self.logit_scale = logit_scale\n",
    "        self.modality1_features = modality1_features\n",
    "        self.modality2_features = modality2_features\n",
    "        self.process_precalculated_features(dataset_name)\n",
    "\n",
    "    def process_precalculated_features(self, dataset_name):\n",
    "        data_prefix = dataset_name + '_' + self.model_name + '_' + self.name\n",
    "        data_prefix = data_prefix.replace('/','-')\n",
    "        np.savetxt(am_utils.data_checkpoint_dir + data_prefix + '_image-embedding.csv', self.modality1_features.cpu(), delimiter = ',')\n",
    "        np.savetxt(am_utils.data_checkpoint_dir + data_prefix + '_text-embedding.csv', self.modality2_features.cpu(), delimiter = ',') \n",
    "\n",
    "    def encode_image(self, images):\n",
    "        raise NotImplementedError(\"this cannot be done for precalculated features -> use cached features\")\n",
    "    \n",
    "    def encode_text(self, texts):\n",
    "        raise NotImplementedError(\"this cannot be done for precalculated features -> use cached features\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "molecule_features = os.path.join(datapath, \"all_molecule_cellpainting_features.pkl\")\n",
    "image_features = os.path.join(datapath, \"subset_image_cellpainting_features.pkl\")\n",
    "\n",
    "# molecule features\n",
    "mol_features_torch = torch.load(molecule_features, map_location=device)\n",
    "mol_features = mol_features_torch[\"mol_features\"]\n",
    "mol_ids = mol_features_torch[\"mol_ids\"]\n",
    "\n",
    "# microscopy features\n",
    "img_features_torch = torch.load(image_features, map_location=device)\n",
    "img_features = img_features_torch[\"img_features\"]\n",
    "img_ids = img_features_torch[\"img_ids\"]\n",
    "\n",
    "# extract subsets of features\n",
    "img_feature_idcs = [np.where(np.array(img_features_torch[\"img_ids\"])==i)[0][0] for i in dataset_cloome.dataset[\"SAMPLE_KEY_img\"].values]\n",
    "mol_feature_idcs = [np.where(np.array(mol_features_torch[\"mol_ids\"])==i)[0][0] for i in dataset_cloome.dataset[\"SAMPLE_KEY_mol\"].values]\n",
    "\n",
    "mol_features_sample = mol_features_torch['mol_features'][mol_feature_idcs]\n",
    "mol_features_sample = am_utils.l2_norm(mol_features_sample)\n",
    "\n",
    "img_features_sample = img_features_torch['img_features'][img_feature_idcs]\n",
    "img_features_sample = am_utils.l2_norm(img_features_sample)\n",
    "\n",
    "# cache features\n",
    "model = PrecalculatedModel('seed-%i'%seed, cloome_dataset_name, img_features_sample, mol_features_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "found cached embeddings for CLOOMDataset_size-100_precalculated_seed-31415\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "87d1d581c2c94464bee9d98a6aea0579",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "CLIPExplorerWidget(children=(VBox(children=(HBox(children=(Dropdown(description='Model: ', options=('precalcul…"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cloome_widget = am_widgets.CLIPExplorerWidget(cloome_dataset_name, cloome_images, cloome_molecules, models=[model])\n",
    "cloome_widget.hover_widget.width = 200\n",
    "cloome_widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cloome",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}