Update app

app.py

@@ -1,7 +1,6 @@
-import os
+import io
 import os.path
 import sys
-from os.path import splitext
 
 import gradio as gr
 import matplotlib.pyplot as plt
@@ -19,6 +18,7 @@ from skimage.color import label2rgb
 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):
@@ -67,6 +67,8 @@ def get_diagonal(W: scipy.sparse.csr_matrix, threshold: float = 1e-12):
 model_name = 'dino_vitb16'  # TODOL Figure out how to make this user-editable
 K = 5
 
+# Fixed parameters
+MAX_SIZE = 384
 
 # Load model
 model, val_transform, patch_size, num_heads = get_model(model_name)
@@ -122,7 +124,7 @@ def segment(inp: Image):
 
     # Remove hook from the model
     handle.remove()
-
+    
     # Normalize features
     normalize = True
     if normalize:
@@ -160,27 +162,36 @@ def segment(inp: Image):
             eigenvectors[k] = 0 - eigenvectors[k]
 
     # Arrange eigenvectors into grid
-    output_image_grid = []
-    for i in range(1, K):
-        eigenvector = eigenvectors[i].reshape(1, 1, H_pad, W_pad)
-        eigenvector = F.interpolate(eigenvector, size=(H, W), mode='nearest')  # slightly off, but for visualizations this is okay
-        # plt.imsave('./tmp.png', eigenvector.squeeze().numpy())  # save to a temporary location
-        # eigenvector = Image.open('./tmp.png').convert('RGB') # load back from our temporary location
-        output_image_grid.append(eigenvector)
-    img_tensor_grid = make_grid(output_image_grid, nrow=8, pad_value=1)
-
-    # Postprocess for Gradio
-    img_tensor_grid.numpy().squeeze()
-    
-    return img_tensor_grid
+    cmap = get_cmap('viridis')
+    output_images = []
+    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
+        buffer = io.BytesIO()
+        plt.imsave(buffer, eigenvector.squeeze().numpy(), format='png')  # save to a temporary location
+        buffer.seek(0)
+        eigenvector_vis = Image.open(buffer).convert('RGB')
+        # eigenvector_vis = TF.to_tensor(eigenvector_vis).unsqueeze(0)
+        eigenvector_vis = np.array(eigenvector_vis)
+        output_images.append(eigenvector_vis)
+    # output_images = torch.cat(output_images, dim=0)
+    # output_images = make_grid(output_images, nrow=8, pad_value=1)
+
+    # # Postprocess for Gradio
+    # output_images = np.array(TF.to_pil_image(output_images))
+    print(f'{len(output_images)=}')
+    return output_images
 
 # Placeholders
-input_placeholders = GradioInputImage(shape=(256, 256), source="upload", tool="editor", type="pil")
-output_placeholders = GradioOutputImage(type="numpy", label=f"Eigenvectors")
-# alternatively: [GradioOutputImage(type="numpy", label=f"Eigenvector {i}") for i in range(K)]
+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)]
 
 # Metadata
-examples = [["images/img1.jpg"], ["images/img2.jpg"]]
+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..."
 thumbnail = "https://raw.githubusercontent.com/gradio-app/hub-echonet/master/thumbnail.png"