Update with description and CRF

app.py

@@ -1,4 +1,6 @@
+from collections import namedtuple
 import io
+from typing import Tuple
 
 import gradio as gr
 import matplotlib.pyplot as plt
@@ -9,12 +11,12 @@ import torch.nn.functional as F
 import torchvision.transforms.functional as TF
 from gradio.inputs import Image as GradioInputImage
 from gradio.outputs import Image as GradioOutputImage
+from matplotlib.pyplot import get_cmap
 from PIL import Image
 from scipy.sparse.linalg import eigsh
 from torch.utils.hooks import RemovableHandle
 from torchvision import transforms
 from torchvision.utils import make_grid
-from matplotlib.pyplot import get_cmap
 
 
 def get_model(name: str):
@@ -60,7 +62,7 @@ def get_diagonal(W: scipy.sparse.csr_matrix, threshold: float = 1e-12):
 
 
 # Parameters
-model_name = 'dino_vitb16'  # TODOL Figure out how to make this user-editable
+model_name = 'dino_vitb16'  # TODO: Figure out how to make this user-editable
 K = 5
 
 # Fixed parameters
@@ -160,6 +162,7 @@ def segment(inp: Image):
     # Arrange eigenvectors into grid
     cmap = get_cmap('viridis')
     output_images = []
+    eigenvectors_upscaled = []
     for i in range(1, K + 1):
         eigenvector = eigenvectors[i].reshape(1, 1, H_patch, W_patch)  # .reshape(1, 1, H_pad, W_pad)
         eigenvector: torch.Tensor = F.interpolate(eigenvector, size=(H_pad, W_pad), mode='bilinear', align_corners=False)  # slightly off, but for visualizations this is okay
@@ -169,10 +172,40 @@ def segment(inp: Image):
         eigenvector_vis = Image.open(buffer).convert('RGB')
         # eigenvector_vis = TF.to_tensor(eigenvector_vis).unsqueeze(0)
         eigenvector_vis = np.array(eigenvector_vis)
+        eigenvectors_upscaled.append(eigenvector)
         output_images.append(eigenvector_vis)
     # output_images = torch.cat(output_images, dim=0)
     # output_images = make_grid(output_images, nrow=8, pad_value=1)
 
+    # Also add CRF
+    if False:
+
+        # Imports
+        import denseCRF
+
+        # Parameters
+        ParamsCRF = namedtuple('ParamsCRF', 'w1 alpha beta w2 gamma it')
+        DEFAULT_CRF_PARAMS = ParamsCRF(
+            w1    = 6,     # weight of bilateral term  # 10.0,
+            alpha = 40,    # spatial std  # 80,  
+            beta  = 13,    # rgb  std  # 13,  
+            w2    = 3,     # weight of spatial term  # 3.0, 
+            gamma = 3,     # spatial std  # 3,   
+            it    = 5.0,   # iteration  # 5.0, 
+        )
+
+        # Get unary potentials
+        unary_potentials = eigenvectors_upscaled[0].squeeze(1).squeeze(0)
+        unary_potentials = (unary_potentials - unary_potentials.min()) / (unary_potentials.max() - unary_potentials.min())
+        unary_potentials_np = torch.stack((1 - unary_potentials, unary_potentials), dim=-1).cpu().numpy()
+        img_np = images.cpu().numpy().transpose(0, 2, 3, 1)
+        img_np = (img_np * 255).astype(np.uint8)[0]
+
+        # Return result of CRF
+        out = denseCRF.densecrf(img_np, unary_potentials_np, DEFAULT_CRF_PARAMS)
+        out = out * 255
+        output_images.append(out)
+
     # # Postprocess for Gradio
     # output_images = np.array(TF.to_pil_image(output_images))
     print(f'{len(output_images)=}')
@@ -181,15 +214,23 @@ def segment(inp: Image):
 # Placeholders
 input_placeholders = GradioInputImage(source="upload", tool="editor", type="pil")
 # output_placeholders = GradioOutputImage(type="numpy", label=f"Eigenvectors")
-output_placeholders = [GradioOutputImage(type="numpy", label=f"Eigenvector {i}") for i in range(K)]
+output_placeholders = [GradioOutputImage(type="numpy", label=(f"Eigenvector {i}")) for i in range(K)]
 
 # Metadata
 examples = [f"examples/{stem}.jpg" for stem in [
     '2008_000099', '2008_000499', '2007_009446', '2007_001586', '2010_001256', '2008_000764', '2008_000705',  # '2007_000039'
 ]]
 
-title = "Deep Spectral Segmentation"
-description = "Deep spectral segmentation (add description here) ..."
+title = "Demo: Deep Spectral Methods for Unsupervised Localization and Segmentation"
+description = """
+This is a demo of <a href="https://lukemelas.github.io/deep-spectral-segmentation/">Deep Spectral Methods: A Surprisingly Strong Baseline for Unsupervised Semantic Segmentation and Localization
+</a> (CVPR 2022 Oral).
+
+Our method decomposes an image into a set of soft segments in a <em>completely unsupervised</em> manner. Specifically, we extract the Laplacian eigenvectors of a feature affinity matrix from a large self-supervised network, and we find that we find that these eigenvectors can be readily used to localize and segment objects.
+
+Below, you can upload an image (or select one of the examples) and see how our method decomposes it into soft segments. Hopefully it will localize some of the objects or semantic segments in your image! 
+"""
+
 thumbnail = "https://raw.githubusercontent.com/gradio-app/hub-echonet/master/thumbnail.png"
 
 # Gradio

requirements.txt

@@ -3,4 +3,5 @@ numpy
 matplotlib
 torch==1.9.1+cpu 
 torchvision==0.10.1+cpu
-scipy
+scipy
+# SimpleCRF