init

.gitignore

@@ -0,0 +1,135 @@
+# MacOS & PyCharm
+.DS_store
+.DS_Store
+.idea/
+sync.sh
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

app.py

@@ -0,0 +1,67 @@
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+
+from infer_model import CLIPpyModel
+from utils import get_similarity, get_transform, ade_palette, get_cmap_image
+
+pretrained_ckpt = "https://github.com/kahnchana/clippy/releases/download/v1.0/clippy_5k.pt"
+ckpt = torch.utils.model_zoo.load_url(pretrained_ckpt)
+
+clippy = CLIPpyModel()
+transform = get_transform((224, 224))
+
+msg = clippy.load_state_dict(ckpt, strict=False)
+
+palette = ade_palette()
+
+
+def process_image(img, captions):
+    sample_text = [x.strip() for x in captions.split(",")]
+    sample_prompts = [f"a photo of a {x}" for x in sample_text]
+
+    image = Image.fromarray(img)
+    image_vector = clippy.encode_image(transform(image).unsqueeze(0), get_pos_tokens=True)
+    text_vector = clippy.text.encode(sample_prompts, convert_to_tensor=True)
+
+    similarity = get_similarity(image_vector, text_vector, (224, 224), do_argmax=True)[0, 0].numpy()
+    rgb_seg = np.zeros((similarity.shape[0], similarity.shape[1], 3), dtype=np.uint8)
+    for idx, _ in enumerate(sample_text):
+        rgb_seg[similarity == idx] = palette[idx]
+
+    joint = Image.blend(image, Image.fromarray(rgb_seg), 0.5)
+    cmap = get_cmap_image({label: tuple(palette[idx]) for idx, label in enumerate(sample_text)})
+
+    return cmap, rgb_seg, joint
+
+
+title = 'CLIPpy'
+
+description = """
+Gradio Demo for CLIPpy: Perceptual Grouping in Contrastive Vision Language Models. \n \n
+Upload an image and type in a set of comma separated labels (e.g.: "man, woman, background"). 
+CLIPPy will segment the image, according to the set of class label you provide.
+"""
+
+article = """
+<p style='text-align: center'>
+<a href='https://arxiv.org/abs/2210.09996' target='_blank'>
+Perceptual Grouping in Contrastive Vision Language Models
+</a>
+|
+<a href='https://github.com/kahnchana/clippy' target='_blank'>Github Repository</a></p>
+"""
+
+demo = gr.Interface(
+    fn=process_image,
+    inputs=[gr.Image(shape=(224, 224)), "text"],
+    outputs=[gr.Image(shape=(224, 224)).style(height=150),
+             gr.Image(shape=(224, 224)).style(height=260),
+             gr.Image(shape=(224, 224)).style(height=260)],
+    title=title,
+    description=description,
+    article=article,
+)
+
+demo.launch()

infer_model.py

@@ -0,0 +1,202 @@
+import abc
+import math
+
+import torch
+import torch.nn.functional as F
+from sentence_transformers import SentenceTransformer
+from timm.models.vision_transformer import (
+    VisionTransformer,
+    build_model_with_cfg,
+    checkpoint_filter_fn,
+    checkpoint_seq,
+    resolve_pretrained_cfg,
+)
+from torch import Tensor, nn
+
+
+class BlankLayer(nn.Module):
+    pass
+
+
+class CustomViT(VisionTransformer):
+    def __init__(
+            self,
+            *args,
+            image_pooling="gmp",
+            **kwargs,
+    ):
+        super(CustomViT, self).__init__(
+            *args, **kwargs
+        )
+        self.image_pooling = image_pooling
+
+    def forward_head(self, x, pre_logits: bool = False):
+        if self.image_pooling:
+            if self.image_pooling == "gap":
+                x = x[:, self.num_prefix_tokens:].mean(dim=1)
+            elif self.image_pooling == "gmp":
+                x = x[:, self.num_prefix_tokens:].max(dim=-2)[0]
+            elif self.image_pooling == "all":
+                x = x[:, self.num_prefix_tokens:]
+            else:  # cls by default
+                x = x[:, 0]
+        x = self.fc_norm(x)
+        return x if pre_logits else self.head(x)
+
+    def forward(self, x, get_pos_tokens=False):
+        x = self.forward_features(x, get_pos_tokens=get_pos_tokens)
+        if get_pos_tokens:
+            return self.fc_norm(x[:, self.num_prefix_tokens:])
+        x = self.forward_head(x)
+        return x
+
+    def forward_features(self, x, get_pos_tokens=False):
+        _, nc, h, w = x.shape
+        x = self.patch_embed(x)
+        x = self._pos_embed(x, w, h)
+        if self.grad_checkpointing and not torch.jit.is_scripting():
+            x = checkpoint_seq(self.blocks, x)
+        else:
+            x = self.blocks(x)
+        x = self.norm(x)
+        return x
+
+    def _pos_embed(self, x, w, h):
+        if self.no_embed_class:
+            # deit-3, updated JAX (big vision)
+            # position embedding does not overlap with class token, add then concat
+            x = x + self.pos_embed
+            if self.cls_token is not None:
+                x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+        else:
+            # original timm, JAX, and deit vit impl
+            # pos_embed has entry for class token, concat then add
+            if self.cls_token is not None:
+                x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+            x = x + self._interpolate_pos_encoding(x, w, h)
+        return self.pos_drop(x)
+
+    def _interpolate_pos_encoding(self, x, w, h):
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        class_pos_embed = self.pos_embed[:, 0]
+        patch_pos_embed = self.pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_embed.patch_size[0]
+        h0 = h // self.patch_embed.patch_size[1]
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w0 + 0.1, h0 + 0.1
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(
+                0, 3, 1, 2
+            ),
+            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
+            mode="bicubic",
+        )
+        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+
+def _create_vision_transformer(variant, pretrained=False, **kwargs):
+    if kwargs.get("features_only", None):
+        raise RuntimeError("features_only not implemented for Vision Transformer models.")
+
+    pretrained_cfg = resolve_pretrained_cfg(
+        variant, pretrained_cfg=kwargs.pop("pretrained_cfg", None)
+    )
+    model = build_model_with_cfg(
+        CustomViT,
+        variant,
+        pretrained,
+        pretrained_cfg=pretrained_cfg,
+        pretrained_filter_fn=checkpoint_filter_fn,
+        pretrained_custom_load="npz" in pretrained_cfg["url"],
+        **kwargs,
+    )
+    return model
+
+
+def vit_base_patch16_224(pretrained=False, variant="vit_base_patch16_224_dino", **kwargs):
+    """ViT-Base (ViT-B/16) /w DINO pretrained weights (no head) - https://arxiv.org/abs/2104.14294"""
+    model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)
+    model = _create_vision_transformer(variant, pretrained=pretrained, **model_kwargs)
+    return model
+
+
+class CLIPpyModel(abc.ABC, torch.nn.Module):
+    """ Implements code for running inference with pre-trained CLIPpy model.
+
+    NOTE: weights used are for a model trained with lower batch-size leading to results below those in paper.
+    """
+
+    def __init__(
+            self,
+            image_pooling: str = "cls",
+            text_pooling: str = "gap",
+    ):
+        super().__init__()
+
+        self.visual = BlankLayer()
+
+        self.visual.trunk = vit_base_patch16_224(True, image_pooling=image_pooling)
+
+        self.text = SentenceTransformer("sentence-transformers/sentence-t5-base")
+        self.logit_scale = nn.Parameter(torch.ones([]) * math.log(1 / 0.07))
+        self.set_text_pooling(text_pooling)
+
+        self._divisor_eps = 1e-4
+        self._image_pooling = image_pooling
+        self._text_pooling = text_pooling
+
+    def forward(
+            self,
+            images: Tensor,
+            input_ids: Tensor,
+            input_id_masks: Tensor,
+            get_pos_tokens: bool = False,
+            **kwargs,
+    ):
+
+        image_encodings = self.encode_image(images, get_pos_tokens=get_pos_tokens)
+
+        if get_pos_tokens:
+            return {
+                image_encodings: image_encodings,
+            }
+
+        text_encodings = self.encode_text(input_ids, input_id_masks)
+
+        return {
+            image_encodings: image_encodings,
+            text_encodings: text_encodings,
+        }
+
+    def encode_text(self, input_ids: Tensor, input_id_masks: Tensor = None, **kwargs):
+        output = self.text({"input_ids": input_ids, "attention_mask": input_id_masks})[
+            "sentence_embedding"
+        ]
+        return self.text_head(output)
+
+    def text_head(self, hidden_states: Tensor, input_id_masks: Tensor = None, **kwargs):
+        return F.normalize(hidden_states, dim=-1, eps=self._divisor_eps).float()
+
+    def encode_image(self, images: Tensor, get_pos_tokens: bool = False, **kwargs):
+        output = self.visual.trunk(images, get_pos_tokens)
+        return self.image_head(output, get_pos_tokens=get_pos_tokens)
+
+    def image_head(self, hidden_states: Tensor, get_pos_tokens: bool = False, **kwargs):
+        return F.normalize(hidden_states, dim=-1, eps=self._divisor_eps).float()
+
+    def set_text_pooling(self, pooling):
+        """ Converts pooling in the Hugging Face model to be max or average pooling"""
+        if pooling == "gmp":
+            self.text[1].pooling_mode_mean_tokens = False
+            self.text[1].pooling_mode_max_tokens = True
+        elif pooling == "gap":
+            pass
+        else:
+            raise NotImplementedError(f"{pooling} not implemented")

utils.py

@@ -0,0 +1,166 @@
+import matplotlib
+import matplotlib.cm as cm
+import matplotlib.colors as mcolors
+import numpy as np
+import torch
+import torchvision
+from PIL import Image, ImageDraw, ImageFont
+from einops import rearrange
+from matplotlib import pyplot as plt
+
+
+def get_similarity(image_encodings, label_encodings, target_shape, interpolation="bilinear", do_argmax=False):
+    """
+
+    Args:
+        image_encodings:
+        label_encodings:
+        target_shape:
+        interpolation: nearest, bilinear
+        do_argmax:
+
+    Returns:
+
+    """
+
+    image_encodings = image_encodings.cpu()
+    label_encodings = label_encodings.cpu()
+
+    image_encodings = rearrange(
+        image_encodings, "b (h w) d -> d b h w", h=int(np.sqrt(image_encodings.shape[-2]))
+    )
+    # assuming square inputs & targets
+    scale_ratio = (target_shape[-2] / image_encodings.shape[-2],
+                   target_shape[-1] / image_encodings.shape[-1],)
+    temp_list = []
+    for i in image_encodings:
+        i = i.unsqueeze(1)
+        i = torch.nn.functional.interpolate(
+            i, scale_factor=scale_ratio, mode=interpolation
+        )
+        temp_list.append(i)
+    image_encodings = torch.cat(temp_list, dim=1)
+
+    image_encodings = rearrange(image_encodings, "b d h w -> b h w d")
+    similarity = image_encodings @ label_encodings.T
+    similarity = rearrange(similarity, "b h w d-> b d h w")
+    if do_argmax:
+        similarity = torch.argmax(similarity, dim=1, keepdim=True).to(torch.float64)
+    return similarity
+
+
+def get_cmap(ncolors):
+    if ncolors > 9:
+        cmap = plt.cm.tab20
+    else:
+        cmap = plt.cm.tab10
+    cmaplist = [cmap(i) for i in range(ncolors)]
+    cmap = matplotlib.colors.LinearSegmentedColormap.from_list("custom", cmaplist, ncolors)
+
+    mappable = cm.ScalarMappable(cmap=cmap)
+    mappable.set_array([])
+    mappable.set_clim(-0.5, ncolors + 0.5)
+
+    return cmap, mappable
+
+
+def vis_prediction(sample_text, img_arr, similarity):
+    N = len(sample_text)
+    cmap, mappable = get_cmap(N)
+
+    fig, axs = plt.subplots(1, 2)
+
+    _ = axs[0].imshow(img_arr)
+    _ = axs[1].imshow(img_arr)
+    _ = axs[1].imshow(similarity, cmap=cmap, interpolation="nearest", vmin=0, vmax=N, alpha=0.5)
+    axs[0].axis("off")
+    axs[1].axis("off")
+
+    fig.subplots_adjust(bottom=0.2)
+    cbar_ax = fig.add_axes([0.0, 0.85, 1.0, 0.05])
+    colorbar = plt.colorbar(mappable, cax=cbar_ax, cmap=cmap, orientation="horizontal")
+    colorbar.set_ticks(np.linspace(0, N, N))
+    colorbar.set_ticklabels(sample_text)
+
+    return fig
+
+
+class DummyArgs:
+    def __init__(self, **kwargs):
+        self.__dict__.update(kwargs)
+
+
+def get_transform(size=(224, 224)):
+    transform = torchvision.transforms.Compose([
+        torchvision.transforms.Resize(size),
+        torchvision.transforms.ToTensor(),
+        torchvision.transforms.Normalize(mean=(0.48145466, 0.4578275, 0.40821073),
+                                         std=(0.26862954, 0.26130258, 0.27577711))
+    ])
+    return transform
+
+
+def ade_palette():
+    """ADE20K palette that maps each class to RGB values."""
+    return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
+            [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
+            [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
+            [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
+            [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
+            [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
+            [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220],
+            [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224],
+            [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255],
+            [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7],
+            [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],
+            [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255],
+            [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0],
+            [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255],
+            [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255],
+            [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],
+            [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0],
+            [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0],
+            [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255],
+            [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255],
+            [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],
+            [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255],
+            [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255],
+            [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255],
+            [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0],
+            [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],
+            [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255],
+            [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112],
+            [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160],
+            [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163],
+            [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],
+            [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0],
+            [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255],
+            [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204],
+            [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255],
+            [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],
+            [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194],
+            [102, 255, 0], [92, 0, 255]]
+
+
+def get_cmap_image(legend):
+    # Define the size of the legend image
+    width = 200
+    height = len(legend) * 20
+
+    # Create a new image with the desired size and background color
+    img = Image.new('RGB', (width, height), (255, 255, 255))
+
+    # Create a drawing context
+    draw = ImageDraw.Draw(img)
+
+    # Define the font to use for the legend labels
+    font = ImageFont.truetype('arial.ttf', 16)
+
+    # Loop through the items in legend and draw a rectangle and label for each
+    y = 0
+    for label, color in legend.items():
+        draw.rectangle((0, y, 20, y + 20), fill=color)
+        draw.text((30, y), label, font=font, fill=(0, 0, 0))
+        y += 20
+
+    return img