Merge pull request #148 from ViCCo-Group/integrate_dreamsim

Integrate dreamsim

README.md

@@ -111,6 +111,14 @@ $ pip install git+https://github.com/serre-lab/Harmonization.git
 $ pip install keras-cv-attention-models>=1.3.5
 ```
 
+If you want to extract features for [DreamSim](https://dreamsim-nights.github.io/) from the [DreamSim repo](https://github.com/ssundaram21/dreamsim), you have to additionally run the following `pip` command in your `thingsvision` environment,
+
+```bash
+$ pip install dreamsim==0.1.2
+```
+
+See the [docs](https://thingsvision.github.io/AvailableModels.html) for which `DreamSim` models are available in `thingsvision`.
+
 #### Google Colab.
 Alternatively, you can use Google Colab to play around with `thingsvision` by uploading your image data to Google Drive (via directory mounting).
 You can find the jupyter notebook using `PyTorch` [here](https://colab.research.google.com/github/ViCCo-Group/thingsvision/blob/master/notebooks/pytorch.ipynb) and the `TensorFlow` example [here](https://colab.research.google.com/github/ViCCo-Group/thingsvision/blob/master/notebooks/tensorflow.ipynb).

docs/AvailableModels.md

@@ -273,3 +273,30 @@ extractor = get_extractor(
   model_parameters=model_parameters
 )
 ```
+
+### DreamSim
+We provide the [DreamSim](https://dreamsim-nights.github.io/) model from the original [DreamSim repo](https://github.com/ssundaram21/dreamsim). To extract features, first install the `dreamsim` package with the following `pip` command 
+```bash
+ $ pip install dreamsim==0.1.2
+ ```
+The model name is:
+- `DreamSim`
+
+We provide two `DreamSim` architectures: CLIP ViT-B/32 and OpenCLIP ViT-B/32. This can be specified using the `model_parameters` argument. For instance, to get the OpenCLIP variant of DreamSim you would do the following:
+```python
+model_name = 'DreamSim'
+source = 'custom'
+device = 'cuda'
+model_parameters = {
+    'variant': 'open_clip_vitb32'
+}
+
+extractor = get_extractor(
+  model_name=model_name,
+  source=source,
+  device=device,
+  pretrained=True,
+  model_parameters=model_parameters
+)
+```
+To load the CLIP ViT-B/32 variant, pass `'clip_vitb32'` to the `variant` parameter instead.

notebooks/pytorch.ipynb

@@ -1,6 +1,7 @@
 {
  "cells": [
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -9,13 +10,15 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "## Settings"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "3ICWd-3iA671"
@@ -67,13 +70,15 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "### Image and feature directories"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "PHSuNkaIAZw2"
@@ -94,6 +99,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "xF0R7sFu-7gI"
@@ -129,6 +135,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -211,14 +218,15 @@
     "                class_names=class_names,\n",
     "                file_names=file_names,\n",
     "            )\n",
-    "            # replace with e.g., [f'conv_{l:02d}'] or [f'fc_{l:02d}']\n",
-    "            features_per_layer[f'layer_{l:02d}'] = features\n",
+    "            # if you want, replace for simplicity with e.g., [f'conv_{l:02d}'], [f'fc_{l:02d}'], or [f'layer_{l:02d}']\n",
+    "            features_per_layer[f'{module_name}'] = features\n",
     "            # save features to disk\n",
     "            save_features(features, out_path=f'{out_path}/features_{model_name}_{module_name}', file_format=file_format)\n",
     "    return features_per_layer"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -245,20 +253,23 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "Select `model` and `layer` for which you want to extract image features. If you want to extract features from a `torchvision` model, use the model naming defined [here](https://pytorch.org/vision/stable/models.html) (e.g., `vgg16` if you want to use VGG-16). If you are uncertain about the naming and enumeration of the layers, use `extractor.show_model()` to see how specific layers are called. This command will show the architecture of the model."
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "### Example 1: VGG-16 with batch norm (pretrained on ImageNet)"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -300,6 +311,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -335,6 +347,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -388,6 +401,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -430,6 +444,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -463,6 +478,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -516,6 +532,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -558,6 +575,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -591,6 +609,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -634,6 +653,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [

notebooks/tensorflow.ipynb

@@ -1,6 +1,7 @@
 {
  "cells": [
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -9,13 +10,15 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "## Settings"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "3ICWd-3iA671"
@@ -66,13 +69,15 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "### Image and feature directories"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "PHSuNkaIAZw2"
@@ -93,6 +98,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "xF0R7sFu-7gI"
@@ -128,6 +134,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -214,14 +221,15 @@
     "                                        class_names=class_names,\n",
     "                                        file_names=file_names,\n",
     "            )\n",
-    "            # replace with e.g., [f'conv_{l:02d}'] or [f'fc_{l:02d}']\n",
-    "            features_per_layer[f'{layer}_{l:02d}'] = features\n",
+    "            # if you want, replace for simplicity with e.g., [f'conv_{l:02d}'], [f'fc_{l:02d}'], or [f'{layer}_{l:02d}']\n",
+    "            features_per_layer[f'{module_name}'] = features\n",
     "            # save features to disk\n",
     "            save_features(features, out_path=f'{out_path}/features_{model_name}_{module_name}', file_format='npy')\n",
     "    return features_per_layer"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -247,20 +255,23 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "Select `model` and `layer` for which you want to extract image features. If you want to extract features from a `torchvision` model, use the model naming defined [here](https://pytorch.org/vision/stable/models.html) (e.g., `vgg16` if you want to use VGG-16). If you are uncertain about the naming and enumeration of the layers, use `model.show()` to see how specific layers called."
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "### VGG-16 (pretrained on ImageNet)"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -304,6 +315,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -334,6 +346,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -385,6 +398,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -428,6 +442,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [

requirements.txt

@@ -17,5 +17,6 @@ torch==1.13.*
 torchvision==0.14.*
 torchtyping
 tqdm
+dreamsim==0.1.2
 git+https://github.com/openai/CLIP.git
 git+https://github.com/serre-lab/Harmonization.git

tests/helper.py

@@ -138,6 +138,18 @@
         "pretrained": True,
         "source": "custom",
         "kwargs": {"variant": "ViT_B16"},
+    },
+    "DreamSim": {
+        "modules": ["model.mlp"],
+        "pretrained": True,
+        "source": "custom",
+        "kwargs": {"variant": "clip_vitb32"},
+    },
+    "DreamSim": {
+        "modules": ["model.mlp"],
+        "pretrained": True,
+        "source": "custom",
+        "kwargs": {"variant": "open_clip_vitb32"},
     }
 }
 

thingsvision/core/extraction/base.py

@@ -1,15 +1,14 @@
 import abc
-import re
 import os
+import re
 import warnings
 from typing import Callable, Iterator, List, Optional, Union
 
 import numpy as np
+import torch
 from torchtyping import TensorType
 from tqdm.auto import tqdm
 
-import torch
-
 Array = np.ndarray