chore: format and resize

app.py

@@ -1,13 +1,15 @@
 import os
 import warnings
-import torch
+
+import gradio as gr
 import numpy as np
-from PIL import Image
-from torchvision import transforms as T
+import torch
 import torch.nn.functional as F
-import gradio as gr
-from hydra import compose, initialize
 from huggingface_hub import Repository
+from hydra import compose, initialize
+from PIL import Image
+from torchvision import transforms as T
+
 from models.builder import build_model
 from segmentation.datasets import PascalVOCDataset
 from visualization import mask2rgb
@@ -41,12 +43,13 @@ COLORS = [
 # Initialize Hydra
 initialize(config_path=CONFIG_PATH, version_base=None)
 
+
 # Configuration and Model Initialization
 def load_model():
     Repository(
         local_dir="clip-dinoiser",
         clone_from="ariG23498/clip-dinoiser",
-        use_auth_token=os.environ.get("token")
+        use_auth_token=os.environ.get("token"),
     )
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -55,27 +58,35 @@ def load_model():
 
     model = build_model(cfg.model, class_names=PascalVOCDataset.CLASSES).to(device)
     model.clip_backbone.decode_head.use_templates = False
-    model.load_state_dict(checkpoint['model_state_dict'], strict=False)
+    model.load_state_dict(checkpoint["model_state_dict"], strict=False)
     return model.eval()
 
+
 def run_clip_dinoiser(input_image, text_prompts, model, device, colors):
-    image = input_image.convert("RGB")
+    # Resize the input image
+    image = input_image.resize((400, 700))
+
+    image = image.convert("RGB")
     text_prompts = text_prompts.split(",")
-    palette = colors[:len(text_prompts)]
+    palette = colors[: len(text_prompts)]
 
     model.clip_backbone.decode_head.update_vocab(text_prompts)
     model.to(device)
 
-    img_tens = T.PILToTensor()(image).unsqueeze(0).to(device) / 255.
+    img_tens = T.PILToTensor()(image).unsqueeze(0).to(device) / 255.0
     h, w = img_tens.shape[-2:]
     output = model(img_tens).cpu()
-    output = F.interpolate(output, scale_factor=model.clip_backbone.backbone.patch_size, mode="bilinear", align_corners=False)[..., :h, :w]
+    output = F.interpolate(
+        output,
+        scale_factor=model.clip_backbone.backbone.patch_size,
+        mode="bilinear",
+        align_corners=False,
+    )[..., :h, :w]
     output = output[0].argmax(dim=0)
 
     mask = mask2rgb(output, palette)
-    classes = np.unique(output).tolist()
     alpha = 0.5
-    blend = (alpha * np.array(image) / 255.) + ((1 - alpha) * mask / 255.)
+    blend = (alpha * np.array(image) / 255.0) + ((1 - alpha) * mask / 255.0)
 
     h_text = [(text, f"{idx}") for idx, text in enumerate(text_prompts)]
     return blend, mask, h_text
@@ -108,22 +119,26 @@ def setup_gradio_interface(model, device, colors, color_map):
                     label="Labels",
                     combine_adjacent=False,
                     show_legend=False,
-                    color_map=color_map
+                    color_map=color_map,
                 )
 
         run_button.click(
-            fn=lambda img, prompts: run_clip_dinoiser(img, prompts, model, device, colors),
+            fn=lambda img, prompts: run_clip_dinoiser(
+                img, prompts, model, device, colors
+            ),
             inputs=[input_image, text_prompts],
-            outputs=[overlay_mask, only_mask, h_text]
+            outputs=[overlay_mask, only_mask, h_text],
         )
 
         gr.Examples(
             examples=[["vintage_bike.jpeg", "background, vintage bike, leather bag"]],
             inputs=[input_image, text_prompts],
             outputs=[overlay_mask, only_mask, h_text],
-            fn=lambda img, prompts: run_clip_dinoiser(img, prompts, model, device, colors),
+            fn=lambda img, prompts: run_clip_dinoiser(
+                img, prompts, model, device, colors
+            ),
             cache_examples=True,
-            label='Try this example input!'
+            label="Try this example input!",
         )
 
     return block