upload clipseg

clipseg/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+This license does not apply to the model weights.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

clipseg/Quickstart.ipynb

@@ -0,0 +1,107 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import requests\n",
+    "\n",
+    "! wget https://owncloud.gwdg.de/index.php/s/ioHbRzFx6th32hn/download -O weights.zip\n",
+    "! unzip -d weights -j weights.zip\n",
+    "from models.clipseg import CLIPDensePredT\n",
+    "from PIL import Image\n",
+    "from torchvision import transforms\n",
+    "from matplotlib import pyplot as plt\n",
+    "\n",
+    "# load model\n",
+    "model = CLIPDensePredT(version='ViT-B/16', reduce_dim=64)\n",
+    "model.eval();\n",
+    "\n",
+    "# non-strict, because we only stored decoder weights (not CLIP weights)\n",
+    "model.load_state_dict(torch.load('weights/rd64-uni.pth', map_location=torch.device('cpu')), strict=False);"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Load and normalize `example_image.jpg`. You can also load through an URL."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# load and normalize image\n",
+    "input_image = Image.open('example_image.jpg')\n",
+    "\n",
+    "# or load from URL...\n",
+    "# image_url = 'https://farm5.staticflickr.com/4141/4856248695_03475782dc_z.jpg'\n",
+    "# input_image = Image.open(requests.get(image_url, stream=True).raw)\n",
+    "\n",
+    "transform = transforms.Compose([\n",
+    "    transforms.ToTensor(),\n",
+    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
+    "    transforms.Resize((352, 352)),\n",
+    "])\n",
+    "img = transform(input_image).unsqueeze(0)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Predict and visualize (this might take a few seconds if running without GPU support)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "prompts = ['a glass', 'something to fill', 'wood', 'a jar']\n",
+    "\n",
+    "# predict\n",
+    "with torch.no_grad():\n",
+    "    preds = model(img.repeat(4,1,1,1), prompts)[0]\n",
+    "\n",
+    "# visualize prediction\n",
+    "_, ax = plt.subplots(1, 5, figsize=(15, 4))\n",
+    "[a.axis('off') for a in ax.flatten()]\n",
+    "ax[0].imshow(input_image)\n",
+    "[ax[i+1].imshow(torch.sigmoid(preds[i][0])) for i in range(4)];\n",
+    "[ax[i+1].text(0, -15, prompts[i]) for i in range(4)];"
+   ]
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "800ed241f7db2bd3aa6942aa3be6809cdb30ee6b0a9e773dfecfa9fef1f4c586"
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "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.8.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

clipseg/Readme.md

@@ -0,0 +1,84 @@
+# Image Segmentation Using Text and Image Prompts
+This repository contains the code used in the paper ["Image Segmentation Using Text and Image Prompts"](https://arxiv.org/abs/2112.10003).
+
+**The Paper has been accepted to CVPR 2022!**
+
+<img src="overview.png" alt="drawing" height="200em"/>
+
+The systems allows to create segmentation models without training based on:
+- An arbitrary text query
+- Or an image with a mask highlighting stuff or an object.
+
+### Quick Start
+
+In the `Quickstart.ipynb` notebook we provide the code for using a pre-trained CLIPSeg model. If you run the notebook locally, make sure you downloaded the `rd64-uni.pth` weights, either manually or via git lfs extension.
+It can also be used interactively using [MyBinder](https://mybinder.org/v2/gh/timojl/clipseg/HEAD?labpath=Quickstart.ipynb)
+(please note that the VM does not use a GPU, thus inference takes a few seconds).
+
+
+### Dependencies
+This code base depends on pytorch, torchvision and clip (`pip install git+https://github.com/openai/CLIP.git`).
+Additional dependencies are hidden for double blind review.
+
+
+### Datasets
+
+* `PhraseCut` and `PhraseCutPlus`: Referring expression dataset
+* `PFEPascalWrapper`: Wrapper class for PFENet's Pascal-5i implementation
+* `PascalZeroShot`: Wrapper class for PascalZeroShot
+* `COCOWrapper`: Wrapper class for COCO.
+
+### Models
+
+* `CLIPDensePredT`: CLIPSeg model with transformer-based decoder.
+* `ViTDensePredT`: CLIPSeg model with transformer-based decoder.
+
+### Third Party Dependencies
+For some of the datasets third party dependencies are required. Run the following commands in the `third_party` folder.  
+```bash
+git clone https://github.com/cvlab-yonsei/JoEm
+git clone https://github.com/Jia-Research-Lab/PFENet.git
+git clone https://github.com/ChenyunWu/PhraseCutDataset.git
+git clone https://github.com/juhongm999/hsnet.git
+```
+
+### Weights
+
+The MIT license does not apply to these weights. 
+
+We provide two model weights, for D=64 (4.1MB) and D=16 (1.1MB).
+```
+wget https://owncloud.gwdg.de/index.php/s/ioHbRzFx6th32hn/download -O weights.zip
+unzip -d weights -j weights.zip
+```
+
+
+### Training and Evaluation
+
+To train use the `training.py` script with experiment file and experiment id parameters. E.g. `python training.py phrasecut.yaml 0` will train the first phrasecut experiment which is defined by the `configuration` and first `individual_configurations` parameters. Model weights will be written in `logs/`.
+
+For evaluation use `score.py`. E.g. `python score.py phrasecut.yaml 0 0` will train the first phrasecut experiment of `test_configuration` and the first configuration in `individual_configurations`.
+
+
+### Usage of PFENet Wrappers
+
+In order to use the dataset and model wrappers for PFENet, the PFENet repository needs to be cloned to the root folder.
+`git clone https://github.com/Jia-Research-Lab/PFENet.git `
+
+
+### License
+
+The source code files in this repository (excluding model weights) are released under MIT license.
+
+### Citation
+```
+@InProceedings{lueddecke22_cvpr,
+    author    = {L\"uddecke, Timo and Ecker, Alexander},
+    title     = {Image Segmentation Using Text and Image Prompts},
+    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+    month     = {June},
+    year      = {2022},
+    pages     = {7086-7096}
+}
+
+```

clipseg/Tables.ipynb

@@ -0,0 +1,349 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "\n",
+    "import clip\n",
+    "from evaluation_utils import norm, denorm\n",
+    "from general_utils import *\n",
+    "from datasets.lvis_oneshot3 import LVIS_OneShot3, LVIS_OneShot"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# PhraseCut"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pc = experiment('experiments/phrasecut.yaml', nums=':6').dataframe()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tab1 = pc[['name', 'pc_miou_best',  'pc_fgiou_best', 'pc_ap']]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cols = ['pc_miou_0.3',  'pc_fgiou_0.3', 'pc_ap']\n",
+    "tab1 = pc[['name'] + cols]\n",
+    "for k in cols:\n",
+    "    tab1.loc[:, k] = (100 * tab1.loc[:, k]).round(1)\n",
+    "tab1.loc[:, 'name'] = ['CLIPSeg (PC+)', 'CLIPSeg (PC, $D=128$)', 'CLIPSeg (PC)', 'CLIP-Deconv', 'ViTSeg (PC+)', 'ViTSeg (PC)']\n",
+    "tab1.insert(1, 't', [0.3]*tab1.shape[0])\n",
+    "print(tab1.to_latex(header=False, index=False))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "For 0.1 threshold"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cols = ['pc_miou_0.1',  'pc_fgiou_0.1', 'pc_ap']\n",
+    "tab1 = pc[['name'] + cols]\n",
+    "for k in cols:\n",
+    "    tab1.loc[:, k] = (100 * tab1.loc[:, k]).round(1)\n",
+    "tab1.loc[:, 'name'] = ['CLIPSeg (PC+)', 'CLIPSeg (PC, $D=128$)', 'CLIPSeg (PC)', 'CLIP-Deconv', 'ViTSeg (PC+)', 'ViTSeg (PC)']\n",
+    "tab1.insert(1, 't', [0.1]*tab1.shape[0])\n",
+    "print(tab1.to_latex(header=False, index=False))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# One-shot"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Pascal"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pas = experiment('experiments/pascal_1shot.yaml', nums=':19').dataframe()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pas[['name', 'pas_h2_miou_0.3', 'pas_h2_biniou_0.3', 'pas_h2_ap', 'pas_h2_fgiou_ct']]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pas = experiment('experiments/pascal_1shot.yaml', nums=':8').dataframe()\n",
+    "tab1 = pas[['pas_h2_miou_0.3', 'pas_h2_biniou_0.3', 'pas_h2_ap']]\n",
+    "print('CLIPSeg (PC+) & 0.3 & CLIP & ' + ' & '.join(f'{x*100:.1f}' for x in tab1[0:4].mean(0).values), '\\\\\\\\')\n",
+    "print('CLIPSeg (PC)  & 0.3 & CLIP & ' + ' & '.join(f'{x*100:.1f}' for x in tab1[4:8].mean(0).values), '\\\\\\\\')\n",
+    "\n",
+    "pas = experiment('experiments/pascal_1shot.yaml', nums='12:16').dataframe()\n",
+    "tab1 = pas[['pas_h2_miou_0.2', 'pas_h2_biniou_0.2', 'pas_h2_ap']]\n",
+    "print('CLIP-Deconv (PC+) & 0.2 & CLIP & ' + ' & '.join(f'{x*100:.1f}' for x in tab1[0:4].mean(0).values), '\\\\\\\\')\n",
+    "\n",
+    "pas = experiment('experiments/pascal_1shot.yaml', nums='16:20').dataframe()\n",
+    "tab1 = pas[['pas_t_miou_0.2', 'pas_t_biniou_0.2', 'pas_t_ap']]\n",
+    "print('ViTSeg (PC+) & 0.2 & CLIP & ' + ' & '.join(f'{x*100:.1f}' for x in tab1[0:4].mean(0).values), '\\\\\\\\')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Pascal Zero-shot (in one-shot setting)\n",
+    "\n",
+    "Using the same setting as one-shot (hence different from the other zero-shot benchmark)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pas = experiment('experiments/pascal_1shot.yaml', nums=':8').dataframe()\n",
+    "tab1 = pas[['pas_t_miou_0.3', 'pas_t_biniou_0.3', 'pas_t_ap']]\n",
+    "print('CLIPSeg (PC+) & 0.3 & CLIP & ' + ' & '.join(f'{x*100:.1f}' for x in tab1[0:4].mean(0).values), '\\\\\\\\')\n",
+    "print('CLIPSeg (PC) & 0.3 & CLIP & ' + ' & '.join(f'{x*100:.1f}' for x in tab1[4:8].mean(0).values), '\\\\\\\\')\n",
+    "\n",
+    "pas = experiment('experiments/pascal_1shot.yaml', nums='12:16').dataframe()\n",
+    "tab1 = pas[['pas_t_miou_0.3', 'pas_t_biniou_0.3', 'pas_t_ap']]\n",
+    "print('CLIP-Deconv (PC+) & 0.3 & CLIP & ' + ' & '.join(f'{x*100:.1f}' for x in tab1[0:4].mean(0).values), '\\\\\\\\')\n",
+    "\n",
+    "pas = experiment('experiments/pascal_1shot.yaml', nums='16:20').dataframe()\n",
+    "tab1 = pas[['pas_t_miou_0.2', 'pas_t_biniou_0.2', 'pas_t_ap']]\n",
+    "print('ViTSeg (PC+) & 0.2 & CLIP & ' + ' & '.join(f'{x*100:.1f}' for x in tab1[0:4].mean(0).values), '\\\\\\\\')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# without fixed thresholds...\n",
+    "\n",
+    "pas = experiment('experiments/pascal_1shot.yaml', nums=':8').dataframe()\n",
+    "tab1 = pas[['pas_t_best_miou', 'pas_t_best_biniou', 'pas_t_ap']]\n",
+    "print('CLIPSeg (PC+) & CLIP & ' + ' & '.join(f'{x*100:.1f}' for x in tab1[0:4].mean(0).values), '\\\\\\\\')\n",
+    "print('CLIPSeg (PC) & CLIP & ' + ' & '.join(f'{x*100:.1f}' for x in tab1[4:8].mean(0).values), '\\\\\\\\')\n",
+    "\n",
+    "pas = experiment('experiments/pascal_1shot.yaml', nums='12:16').dataframe()\n",
+    "tab1 = pas[['pas_t_best_miou', 'pas_t_best_biniou', 'pas_t_ap']]\n",
+    "print('CLIP-Deconv (PC+) & CLIP & ' + ' & '.join(f'{x*100:.1f}' for x in tab1[0:4].mean(0).values), '\\\\\\\\')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### COCO"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "coco = experiment('experiments/coco.yaml', nums=':29').dataframe()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tab1 = coco[['coco_h2_miou_0.1', 'coco_h2_biniou_0.1', 'coco_h2_ap']]\n",
+    "tab2 = coco[['coco_h2_miou_0.2', 'coco_h2_biniou_0.2', 'coco_h2_ap']]\n",
+    "tab3 = coco[['coco_h2_miou_best', 'coco_h2_biniou_best', 'coco_h2_ap']]\n",
+    "print('CLIPSeg (COCO) & 0.1 & CLIP &  ' + ' & '.join(f'{x*100:.1f}' for x in tab1[:4].mean(0).values), '\\\\\\\\')\n",
+    "print('CLIPSeg (COCO+N)  & 0.1 & CLIP &  ' + ' & '.join(f'{x*100:.1f}' for x in tab1[4:8].mean(0).values), '\\\\\\\\')\n",
+    "print('CLIP-Deconv (COCO+N)  & 0.1 & CLIP &  ' + ' & '.join(f'{x*100:.1f}' for x in tab1[12:16].mean(0).values), '\\\\\\\\')\n",
+    "print('ViTSeg (COCO)  & 0.1 & CLIP &  ' + ' & '.join(f'{x*100:.1f}' for x in tab1[8:12].mean(0).values), '\\\\\\\\')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Zero-shot"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "zs = experiment('experiments/pascal_0shot.yaml', nums=':11').dataframe()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "tab1 = zs[['pas_zs_seen', 'pas_zs_unseen']]\n",
+    "print('CLIPSeg (PC+) & CLIP &  ' + ' & '.join(f'{x*100:.1f}' for x in tab1[8:9].values[0].tolist() + tab1[10:11].values[0].tolist()), '\\\\\\\\')\n",
+    "print('CLIP-Deconv & CLIP &  ' + ' & '.join(f'{x*100:.1f}' for x in tab1[2:3].values[0].tolist()  + tab1[3:4].values[0].tolist()), '\\\\\\\\')\n",
+    "print('ViTSeg & ImageNet-1K &  ' + ' & '.join(f'{x*100:.1f}' for x in tab1[4:5].values[0].tolist()  + tab1[5:6].values[0].tolist()), '\\\\\\\\')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Ablation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ablation = experiment('experiments/ablation.yaml', nums=':8').dataframe()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tab1 = ablation[['name', 'pc_miou_best', 'pc_ap', 'pc-vis_miou_best', 'pc-vis_ap']]\n",
+    "for k in ['pc_miou_best', 'pc_ap', 'pc-vis_miou_best', 'pc-vis_ap']:\n",
+    "    tab1.loc[:, k] = (100 * tab1.loc[:, k]).round(1)\n",
+    "tab1.loc[:, 'name'] = ['CLIPSeg', 'no CLIP pre-training', 'no-negatives', '50% negatives', 'no visual', '$D=16$', 'only layer 3', 'highlight mask']"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(tab1.loc[[0,1,4,5,6,7],:].to_latex(header=False, index=False))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(tab1.loc[[0,1,4,5,6,7],:].to_latex(header=False, index=False))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Generalization"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "generalization = experiment('experiments/generalize.yaml').dataframe()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "gen = generalization[['aff_best_fgiou', 'aff_ap', 'ability_best_fgiou', 'ability_ap', 'part_best_fgiou', 'part_ap']].values"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(\n",
+    "    'CLIPSeg (PC+) & ' + ' & '.join(f'{x*100:.1f}' for x in gen[1]) + ' \\\\\\\\ \\n' + \\\n",
+    "    'CLIPSeg (LVIS)  & ' + ' & '.join(f'{x*100:.1f}' for x in gen[0]) + ' \\\\\\\\ \\n' + \\\n",
+    "    'CLIP-Deconv & ' + ' & '.join(f'{x*100:.1f}' for x in gen[2]) + ' \\\\\\\\ \\n' + \\\n",
+    "    'VITSeg & ' + ' & '.join(f'{x*100:.1f}' for x in gen[3]) + ' \\\\\\\\'\n",
+    ")"
+   ]
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "800ed241f7db2bd3aa6942aa3be6809cdb30ee6b0a9e773dfecfa9fef1f4c586"
+  },
+  "kernelspec": {
+   "display_name": "env2",
+   "language": "python",
+   "name": "env2"
+  },
+  "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.8.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

clipseg/Visual_Feature_Engineering.ipynb

@@ -0,0 +1,366 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Systematic"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "\n",
+    "import clip\n",
+    "from evaluation_utils import norm, denorm\n",
+    "from general_utils import *\n",
+    "from datasets.lvis_oneshot3 import LVIS_OneShot3\n",
+    "\n",
+    "clip_device = 'cuda'\n",
+    "clip_model, preprocess = clip.load(\"ViT-B/16\", device=clip_device)\n",
+    "clip_model.eval();\n",
+    "\n",
+    "from models.clipseg import CLIPDensePredTMasked\n",
+    "\n",
+    "clip_mask_model = CLIPDensePredTMasked(version='ViT-B/16').to(clip_device)\n",
+    "clip_mask_model.eval();"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "lvis = LVIS_OneShot3('train_fixed', mask='separate', normalize=True, with_class_label=True, add_bar=False, \n",
+    "                     text_class_labels=True, image_size=352, min_area=0.1,\n",
+    "                     min_frac_s=0.05, min_frac_q=0.05, fix_find_crop=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plot_data(lvis)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from collections import defaultdict\n",
+    "import json\n",
+    "\n",
+    "lvis_raw = json.load(open(expanduser('~/datasets/LVIS/lvis_v1_train.json')))\n",
+    "lvis_val_raw = json.load(open(expanduser('~/datasets/LVIS/lvis_v1_val.json')))\n",
+    "\n",
+    "objects_per_image = defaultdict(lambda : set())\n",
+    "for ann in lvis_raw['annotations']:\n",
+    "    objects_per_image[ann['image_id']].add(ann['category_id'])\n",
+    "    \n",
+    "for ann in lvis_val_raw['annotations']:\n",
+    "    objects_per_image[ann['image_id']].add(ann['category_id'])    \n",
+    "    \n",
+    "objects_per_image = {o: [lvis.category_names[o] for o in v] for o, v in objects_per_image.items()}\n",
+    "\n",
+    "del lvis_raw, lvis_val_raw"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#bs = 32\n",
+    "#batches = [get_batch(lvis, i*bs, (i+1)*bs, cuda=True) for i in range(10)]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from general_utils import get_batch\n",
+    "from functools import partial\n",
+    "from evaluation_utils import img_preprocess\n",
+    "import torch\n",
+    "\n",
+    "def get_similarities(batches_or_dataset, process, mask=lambda x: None, clipmask=False):\n",
+    "\n",
+    "    # base_words = [f'a photo of {x}' for x in ['a person', 'an animal', 'a knife', 'a cup']]\n",
+    "\n",
+    "    all_prompts = []\n",
+    "    \n",
+    "    with torch.no_grad():\n",
+    "        valid_sims = []\n",
+    "        torch.manual_seed(571)\n",
+    "        \n",
+    "        if type(batches_or_dataset) == list:\n",
+    "            loader = batches_or_dataset  # already loaded\n",
+    "            max_iter = float('inf')\n",
+    "        else:\n",
+    "            loader = DataLoader(batches_or_dataset, shuffle=False, batch_size=32)\n",
+    "            max_iter = 50\n",
+    "        \n",
+    "        global batch\n",
+    "        for i_batch, (batch, batch_y) in enumerate(loader):\n",
+    "            \n",
+    "            if i_batch >= max_iter: break\n",
+    "                \n",
+    "            processed_batch = process(batch)\n",
+    "            if type(processed_batch) == dict:\n",
+    "                \n",
+    "                # processed_batch =  {k: v.to(clip_device) for k, v in processed_batch.items()}\n",
+    "                image_features = clip_mask_model.visual_forward(**processed_batch)[0].to(clip_device).half()\n",
+    "            else:\n",
+    "                processed_batch = process(batch).to(clip_device)\n",
+    "                processed_batch = nnf.interpolate(processed_batch, (224, 224), mode='bilinear')\n",
+    "                #image_features = clip_model.encode_image(processed_batch.to(clip_device)) \n",
+    "                image_features = clip_mask_model.visual_forward(processed_batch)[0].to(clip_device).half()\n",
+    "                \n",
+    "            image_features = image_features / image_features.norm(dim=-1, keepdim=True)\n",
+    "            bs = len(batch[0])\n",
+    "            for j in range(bs):\n",
+    "            \n",
+    "                c, _, sid, qid = lvis.sample_ids[bs * i_batch + j]\n",
+    "                support_image = basename(lvis.samples[c][sid])\n",
+    "                \n",
+    "                img_objs = [o for o in objects_per_image[int(support_image)]]\n",
+    "                img_objs = [o.replace('_', ' ') for o in img_objs]\n",
+    "                \n",
+    "                other_words = [f'a photo of a {o.replace(\"_\", \" \")}' for o in img_objs \n",
+    "                               if o != batch_y[2][j]]\n",
+    "            \n",
+    "                prompts = [f'a photo of a {batch_y[2][j]}'] + other_words\n",
+    "                all_prompts += [prompts]\n",
+    "                \n",
+    "                text_cond = clip_model.encode_text(clip.tokenize(prompts).to(clip_device))\n",
+    "                text_cond = text_cond / text_cond.norm(dim=-1, keepdim=True)            \n",
+    "\n",
+    "                global logits\n",
+    "                logits = clip_model.logit_scale.exp() * image_features[j] @ text_cond.T\n",
+    "\n",
+    "                global sim\n",
+    "                sim = torch.softmax(logits, dim=-1)\n",
+    "            \n",
+    "                valid_sims += [sim]\n",
+    "                \n",
+    "        #valid_sims = torch.stack(valid_sims)\n",
+    "        return valid_sims, all_prompts\n",
+    "    \n",
+    "\n",
+    "def new_img_preprocess(x):\n",
+    "    return {'x_inp': x[1], 'mask': (11, 'cls_token', x[2])}\n",
+    "    \n",
+    "#get_similarities(lvis, partial(img_preprocess, center_context=0.5));\n",
+    "get_similarities(lvis, lambda x: x[1]);"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "preprocessing_functions = [\n",
+    "#     ['clip mask CLS L11', lambda x: {'x_inp': x[1].cuda(), 'mask': (11, 'cls_token', x[2].cuda())}],\n",
+    "#     ['clip mask CLS all', lambda x: {'x_inp': x[1].cuda(), 'mask': ('all', 'cls_token', x[2].cuda())}],\n",
+    "#     ['clip mask all all', lambda x: {'x_inp': x[1].cuda(), 'mask': ('all', 'all', x[2].cuda())}],\n",
+    "#     ['colorize object red', partial(img_preprocess, colorize=True)],\n",
+    "#     ['add red outline', partial(img_preprocess, outline=True)],\n",
+    "    \n",
+    "#     ['BG brightness 50%', partial(img_preprocess, bg_fac=0.5)],\n",
+    "#     ['BG brightness 10%', partial(img_preprocess, bg_fac=0.1)],\n",
+    "#     ['BG brightness 0%', partial(img_preprocess, bg_fac=0.0)],\n",
+    "#     ['BG blur', partial(img_preprocess, blur=3)],\n",
+    "#     ['BG blur & intensity 10%', partial(img_preprocess, blur=3, bg_fac=0.1)],\n",
+    "   \n",
+    "#     ['crop large context', partial(img_preprocess, center_context=0.5)],\n",
+    "#     ['crop small context', partial(img_preprocess, center_context=0.1)],\n",
+    "    ['crop & background blur', partial(img_preprocess, blur=3, center_context=0.5)],\n",
+    "    ['crop & intensity 10%', partial(img_preprocess, blur=3, bg_fac=0.1)],\n",
+    "#     ['crop & background blur & intensity 10%', partial(img_preprocess, blur=3, center_context=0.1, bg_fac=0.1)],\n",
+    "]\n",
+    "\n",
+    "preprocessing_functions = preprocessing_functions\n",
+    "\n",
+    "base, base_p = get_similarities(lvis, lambda x: x[1])\n",
+    "outs = [get_similarities(lvis, fun) for _, fun in preprocessing_functions]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "outs2 = [get_similarities(lvis, fun) for _, fun in  [['BG brightness 0%', partial(img_preprocess, bg_fac=0.0)]]]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for j in range(1):\n",
+    "    print(np.mean([outs2[j][0][i][0].cpu() - base[i][0].cpu() for i in range(len(base)) if len(base_p[i]) >= 3]))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pandas import DataFrame\n",
+    "tab = dict()\n",
+    "for j, (name, _) in enumerate(preprocessing_functions):\n",
+    "    tab[name] =  np.mean([outs[j][0][i][0].cpu() - base[i][0].cpu() for i in range(len(base)) if len(base_p[i]) >= 3])\n",
+    "    \n",
+    "    \n",
+    "print('\\n'.join(f'{k} & {v*100:.2f} \\\\\\\\' for k,v in tab.items()))    "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Visual"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from evaluation_utils import denorm, norm"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def load_sample(filename, filename2):\n",
+    "    from os.path import join\n",
+    "    bp = expanduser('~/cloud/resources/sample_images')\n",
+    "    tf = transforms.Compose([\n",
+    "        transforms.ToTensor(),\n",
+    "        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
+    "        transforms.Resize(224),\n",
+    "        transforms.CenterCrop(224)\n",
+    "    ])\n",
+    "    tf2 = transforms.Compose([\n",
+    "        transforms.ToTensor(),\n",
+    "        transforms.Resize(224),\n",
+    "        transforms.CenterCrop(224)\n",
+    "    ])\n",
+    "    inp1 = [None, tf(Image.open(join(bp, filename))), tf2(Image.open(join(bp, filename2)))]\n",
+    "    inp1[1] = inp1[1].unsqueeze(0)\n",
+    "    inp1[2] = inp1[2][:1]   \n",
+    "    return inp1\n",
+    "\n",
+    "def all_preprocessing(inp1):\n",
+    "    return [\n",
+    "        img_preprocess(inp1),\n",
+    "        img_preprocess(inp1, colorize=True),\n",
+    "        img_preprocess(inp1, outline=True),        \n",
+    "        img_preprocess(inp1, blur=3),\n",
+    "        img_preprocess(inp1, bg_fac=0.1),\n",
+    "        #img_preprocess(inp1, bg_fac=0.5),\n",
+    "        #img_preprocess(inp1, blur=3, bg_fac=0.5),        \n",
+    "        img_preprocess(inp1, blur=3, bg_fac=0.5, center_context=0.5),\n",
+    "    ]\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from torchvision import transforms\n",
+    "from PIL import Image\n",
+    "from matplotlib import pyplot as plt\n",
+    "from evaluation_utils import img_preprocess\n",
+    "import clip\n",
+    "\n",
+    "images_queries = [\n",
+    "    [load_sample('things1.jpg', 'things1_jar.png'), ['jug', 'knife', 'car', 'animal', 'sieve', 'nothing']],\n",
+    "    [load_sample('own_photos/IMG_2017s_square.jpg', 'own_photos/IMG_2017s_square_trash_can.png'), ['trash bin', 'house', 'car', 'bike', 'window', 'nothing']],\n",
+    "]\n",
+    "\n",
+    "\n",
+    "_, ax = plt.subplots(2 * len(images_queries), 6, figsize=(14, 4.5 * len(images_queries)))\n",
+    "\n",
+    "for j, (images, objects) in enumerate(images_queries):\n",
+    "    \n",
+    "    joint_image = all_preprocessing(images)\n",
+    "    \n",
+    "    joint_image = torch.stack(joint_image)[:,0]\n",
+    "    clip_model, preprocess = clip.load(\"ViT-B/16\", device='cpu')\n",
+    "    image_features = clip_model.encode_image(joint_image)\n",
+    "    image_features = image_features / image_features.norm(dim=-1, keepdim=True)\n",
+    "    \n",
+    "    prompts = [f'a photo of a {obj}'for obj in objects]\n",
+    "    text_cond = clip_model.encode_text(clip.tokenize(prompts))\n",
+    "    text_cond = text_cond / text_cond.norm(dim=-1, keepdim=True)\n",
+    "    logits = clip_model.logit_scale.exp() * image_features @ text_cond.T\n",
+    "    sim = torch.softmax(logits, dim=-1).detach().cpu()\n",
+    "\n",
+    "    for i, img in enumerate(joint_image):\n",
+    "        ax[2*j, i].axis('off')\n",
+    "        \n",
+    "        ax[2*j, i].imshow(torch.clamp(denorm(joint_image[i]).permute(1,2,0), 0, 1))\n",
+    "        ax[2*j+ 1, i].grid(True)\n",
+    "        \n",
+    "        ax[2*j + 1, i].set_ylim(0,1)\n",
+    "        ax[2*j + 1, i].set_yticklabels([])\n",
+    "        ax[2*j + 1, i].set_xticks([])  # set_xticks(range(len(prompts)))\n",
+    "#         ax[1, i].set_xticklabels(objects, rotation=90)\n",
+    "        for k in range(len(sim[i])):\n",
+    "            ax[2*j + 1, i].bar(k, sim[i][k], color=plt.cm.tab20(1) if k!=0 else plt.cm.tab20(3))\n",
+    "            ax[2*j + 1, i].text(k, 0.07, objects[k], rotation=90, ha='center', fontsize=15)\n",
+    "\n",
+    "plt.tight_layout()\n",
+    "plt.savefig('figures/prompt_engineering.pdf', bbox_inches='tight')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "env2",
+   "language": "python",
+   "name": "env2"
+  },
+  "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.8.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

clipseg/datasets/coco_wrapper.py

@@ -0,0 +1,99 @@
+import pickle
+from types import new_class
+import torch
+import numpy as np
+import os
+import json
+
+from os.path import join, dirname, isdir, isfile, expanduser, realpath, basename
+from random import shuffle, seed as set_seed
+from PIL import Image
+
+from itertools import combinations
+from torchvision import transforms
+from torchvision.transforms.transforms import Resize
+
+from datasets.utils import blend_image_segmentation
+from general_utils import get_from_repository
+
+COCO_CLASSES = {0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', 5: 'bus', 6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light', 10: 'fire hydrant', 11: 'stop sign', 12: 'parking meter', 13: 'bench', 14: 'bird', 15: 'cat', 16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow', 20: 'elephant', 21: 'bear', 22: 'zebra', 23: 'giraffe', 24: 'backpack', 25: 'umbrella', 26: 'handbag', 27: 'tie', 28: 'suitcase', 29: 'frisbee', 30: 'skis', 31: 'snowboard', 32: 'sports ball', 33: 'kite', 34: 'baseball bat', 35: 'baseball glove', 36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 39: 'bottle', 40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife', 44: 'spoon', 45: 'bowl', 46: 'banana', 47: 'apple', 48: 'sandwich', 49: 'orange', 50: 'broccoli', 51: 'carrot', 52: 'hot dog', 53: 'pizza', 54: 'donut', 55: 'cake', 56: 'chair', 57: 'couch', 58: 'potted plant', 59: 'bed', 60: 'dining table', 61: 'toilet', 62: 'tv', 63: 'laptop', 64: 'mouse', 65: 'remote', 66: 'keyboard', 67: 'cell phone', 68: 'microwave', 69: 'oven', 70: 'toaster', 71: 'sink', 72: 'refrigerator', 73: 'book', 74: 'clock', 75: 'vase', 76: 'scissors', 77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'}
+
+class COCOWrapper(object):
+
+    def __init__(self, split, fold=0, image_size=400, aug=None, mask='separate', negative_prob=0,
+                 with_class_label=False):
+        super().__init__()
+
+        self.mask = mask
+        self.with_class_label = with_class_label
+        self.negative_prob = negative_prob
+
+        from third_party.hsnet.data.coco import DatasetCOCO
+
+        get_from_repository('COCO-20i', ['COCO-20i.tar'])
+
+        foldpath = join(dirname(__file__), '../third_party/hsnet/data/splits/coco/%s/fold%d.pkl')
+
+        def build_img_metadata_classwise(self):
+            with open(foldpath % (self.split, self.fold), 'rb') as f:
+                img_metadata_classwise = pickle.load(f)
+            return img_metadata_classwise
+
+
+        DatasetCOCO.build_img_metadata_classwise = build_img_metadata_classwise
+        # DatasetCOCO.read_mask = read_mask
+        
+        mean = [0.485, 0.456, 0.406]
+        std = [0.229, 0.224, 0.225]
+        transform = transforms.Compose([
+            transforms.Resize((image_size, image_size)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean, std)
+        ])
+
+        self.coco = DatasetCOCO(expanduser('~/datasets/COCO-20i/'), fold, transform, split, 1, False)
+    
+        self.all_classes = [self.coco.class_ids]
+        self.coco.base_path = join(expanduser('~/datasets/COCO-20i'))
+
+    def __len__(self):
+        return len(self.coco)
+
+    def __getitem__(self, i):
+        sample = self.coco[i]
+
+        label_name = COCO_CLASSES[int(sample['class_id'])]
+
+        img_s, seg_s = sample['support_imgs'][0], sample['support_masks'][0]
+
+        if self.negative_prob > 0 and torch.rand(1).item() < self.negative_prob:
+            new_class_id = sample['class_id']
+            while new_class_id == sample['class_id']:
+                sample2 = self.coco[torch.randint(0, len(self), (1,)).item()]
+                new_class_id = sample2['class_id']
+            img_s = sample2['support_imgs'][0]
+            seg_s = torch.zeros_like(seg_s)
+
+        mask = self.mask
+        if mask == 'separate':
+            supp = (img_s, seg_s)
+        elif mask == 'text_label':
+            # DEPRECATED
+            supp = [int(sample['class_id'])]
+        elif mask == 'text':
+            supp = [label_name]      
+        else:
+            if mask.startswith('text_and_'):
+                mask = mask[9:]
+                label_add = [label_name]
+            else:
+                label_add = []
+
+            supp = label_add + blend_image_segmentation(img_s, seg_s, mode=mask)
+
+        if self.with_class_label:
+            label = (torch.zeros(0), sample['class_id'],)
+        else:
+            label = (torch.zeros(0), )
+
+        return (sample['query_img'],) + tuple(supp), (sample['query_mask'].unsqueeze(0),) + label

clipseg/datasets/pascal_classes.json

@@ -0,0 +1 @@
+[{"id": 1, "synonyms": ["aeroplane"]}, {"id": 2, "synonyms": ["bicycle"]}, {"id": 3, "synonyms": ["bird"]}, {"id": 4, "synonyms": ["boat"]}, {"id": 5, "synonyms": ["bottle"]}, {"id": 6, "synonyms": ["bus"]}, {"id": 7, "synonyms": ["car"]}, {"id": 8, "synonyms": ["cat"]}, {"id": 9, "synonyms": ["chair"]}, {"id": 10, "synonyms": ["cow"]}, {"id": 11, "synonyms": ["diningtable"]}, {"id": 12, "synonyms": ["dog"]}, {"id": 13, "synonyms": ["horse"]}, {"id": 14, "synonyms": ["motorbike"]}, {"id": 15, "synonyms": ["person"]}, {"id": 16, "synonyms": ["pottedplant"]}, {"id": 17, "synonyms": ["sheep"]}, {"id": 18, "synonyms": ["sofa"]}, {"id": 19, "synonyms": ["train"]}, {"id": 20, "synonyms": ["tvmonitor"]}]

clipseg/datasets/pascal_zeroshot.py

@@ -0,0 +1,60 @@
+from os.path import expanduser
+import torch
+import json
+import torchvision
+from general_utils import get_from_repository
+from general_utils import log
+from torchvision import transforms
+
+PASCAL_VOC_CLASSES_ZS = [['cattle.n.01', 'motorcycle.n.01'], ['aeroplane.n.01', 'sofa.n.01'], 
+                         ['cat.n.01', 'television.n.03'], ['train.n.01', 'bottle.n.01'],
+                          ['chair.n.01', 'pot_plant.n.01']]
+
+
+class PascalZeroShot(object):
+
+    def __init__(self, split, n_unseen, image_size=224) -> None:
+        super().__init__()
+
+        import sys
+        sys.path.append('third_party/JoEm')
+        from third_party.JoEm.data_loader.dataset import VOCSegmentation
+        from third_party.JoEm.data_loader import get_seen_idx, get_unseen_idx, VOC
+
+        self.pascal_classes = VOC
+        self.image_size = image_size
+
+        self.transform = transforms.Compose([
+            transforms.Resize((image_size, image_size)),
+        ])
+
+        if split == 'train':
+            self.voc = VOCSegmentation(get_unseen_idx(n_unseen), get_seen_idx(n_unseen), 
+                                       split=split, transform=True, transform_args=dict(base_size=312, crop_size=312), 
+                                       ignore_bg=False, ignore_unseen=False, remv_unseen_img=True)
+        elif split == 'val':
+            self.voc = VOCSegmentation(get_unseen_idx(n_unseen), get_seen_idx(n_unseen), 
+                                       split=split, transform=False, 
+                                       ignore_bg=False, ignore_unseen=False)
+
+        self.unseen_idx = get_unseen_idx(n_unseen)
+
+    def __len__(self):
+        return len(self.voc)
+
+    def __getitem__(self, i):
+
+        sample = self.voc[i]
+        label = sample['label'].long()
+        all_labels = [l for l in torch.where(torch.bincount(label.flatten())>0)[0].numpy().tolist() if l != 255]
+        class_indices = [l for l in all_labels]
+        class_names = [self.pascal_classes[l] for l in all_labels]
+
+        image = self.transform(sample['image'])
+
+        label = transforms.Resize((self.image_size, self.image_size), 
+            interpolation=torchvision.transforms.InterpolationMode.NEAREST)(label.unsqueeze(0))[0]
+
+        return (image,), (label, )
+
+

clipseg/datasets/pfe_dataset.py

@@ -0,0 +1,129 @@
+from os.path import expanduser
+import torch
+import json
+from general_utils import get_from_repository
+from datasets.lvis_oneshot3 import blend_image_segmentation
+from general_utils import log
+
+PASCAL_CLASSES = {a['id']: a['synonyms'] for a in json.load(open('datasets/pascal_classes.json'))}
+
+
+class PFEPascalWrapper(object):
+
+    def __init__(self, mode, split, mask='separate', image_size=473, label_support=None, size=None, p_negative=0, aug=None):
+        import sys
+        # sys.path.append(expanduser('~/projects/new_one_shot'))
+        from third_party.PFENet.util.dataset import SemData
+
+        get_from_repository('PascalVOC2012', ['Pascal5i.tar'])
+
+        self.p_negative = p_negative
+        self.size = size
+        self.mode = mode
+        self.image_size = image_size
+        
+        if label_support in {True, False}:
+            log.warning('label_support argument is deprecated. Use mask instead.')
+            #raise ValueError()
+
+        self.mask = mask
+
+        value_scale = 255
+        mean = [0.485, 0.456, 0.406]
+        mean = [item * value_scale for item in mean]
+        std = [0.229, 0.224, 0.225]
+        std = [item * value_scale for item in std]
+
+        import third_party.PFENet.util.transform as transform
+
+        if mode == 'val':
+            data_list = expanduser('~/projects/old_one_shot/PFENet/lists/pascal/val.txt')
+
+            data_transform = [transform.test_Resize(size=image_size)] if image_size != 'original' else []
+            data_transform += [
+                transform.ToTensor(),
+                transform.Normalize(mean=mean, std=std)
+            ]   
+
+
+        elif mode == 'train':
+            data_list =  expanduser('~/projects/old_one_shot/PFENet/lists/pascal/voc_sbd_merge_noduplicate.txt')
+
+            assert image_size != 'original'
+
+            data_transform = [
+                transform.RandScale([0.9, 1.1]),
+                transform.RandRotate([-10, 10], padding=mean, ignore_label=255),
+                transform.RandomGaussianBlur(),
+                transform.RandomHorizontalFlip(),
+                transform.Crop((image_size, image_size), crop_type='rand', padding=mean, ignore_label=255),
+                transform.ToTensor(),
+                transform.Normalize(mean=mean, std=std)
+            ]
+
+        data_transform = transform.Compose(data_transform)
+
+        self.dataset = SemData(split=split, mode=mode, data_root=expanduser('~/datasets/PascalVOC2012/VOC2012'), 
+                               data_list=data_list, shot=1, transform=data_transform, use_coco=False, use_split_coco=False)
+
+        self.class_list = self.dataset.sub_val_list if mode == 'val' else self.dataset.sub_list
+
+        # verify that subcls_list always has length 1
+        # assert len(set([len(d[4]) for d in self.dataset])) == 1
+
+        print('actual length', len(self.dataset.data_list))
+
+    def __len__(self):
+        if self.mode == 'val':
+            return len(self.dataset.data_list)
+        else:
+            return len(self.dataset.data_list)
+
+    def __getitem__(self, index):
+        if self.dataset.mode == 'train':
+            image, label, s_x, s_y, subcls_list = self.dataset[index % len(self.dataset.data_list)]
+        elif self.dataset.mode == 'val':
+            image, label, s_x, s_y, subcls_list, ori_label = self.dataset[index % len(self.dataset.data_list)]
+            ori_label = torch.from_numpy(ori_label).unsqueeze(0)
+            
+            if self.image_size != 'original':
+                longerside = max(ori_label.size(1), ori_label.size(2))
+                backmask = torch.ones(ori_label.size(0), longerside, longerside).cuda()*255
+                backmask[0, :ori_label.size(1), :ori_label.size(2)] = ori_label
+                label = backmask.clone().long()      
+            else:
+                label = label.unsqueeze(0) 
+
+            # assert label.shape == (473, 473)
+
+        if self.p_negative > 0:
+            if torch.rand(1).item() < self.p_negative:
+                while True:
+                    idx = torch.randint(0, len(self.dataset.data_list), (1,)).item()
+                    _, _, s_x, s_y, subcls_list_tmp, _ = self.dataset[idx]
+                    if subcls_list[0] != subcls_list_tmp[0]:
+                        break
+
+        s_x = s_x[0]
+        s_y = (s_y == 1)[0]
+        label_fg = (label == 1).float()
+        val_mask = (label != 255).float()
+
+        class_id = self.class_list[subcls_list[0]]
+
+        label_name = PASCAL_CLASSES[class_id][0]
+        label_add = ()
+        mask = self.mask
+
+        if mask == 'text':
+            support = ('a photo of a ' + label_name + '.',)
+        elif mask == 'separate':
+            support = (s_x, s_y)
+        else:
+            if mask.startswith('text_and_'):
+                label_add = (label_name,)
+                mask = mask[9:]
+
+            support = (blend_image_segmentation(s_x, s_y.float(), mask)[0],)
+
+        return (image,) + label_add + support, (label_fg.unsqueeze(0), val_mask.unsqueeze(0), subcls_list[0])