refactor code

app.py

@@ -1,101 +1,96 @@
-from models.builder import build_model
-from visualization import mask2rgb
-from segmentation.datasets import PascalVOCDataset
-
 import os
-from hydra import compose, initialize
+import warnings
+import torch
+import numpy as np
 from PIL import Image
-import matplotlib.pyplot as plt
 from torchvision import transforms as T
 import torch.nn.functional as F
-import numpy as np
-from operator import itemgetter 
-import torch
-import random
-import warnings
-
-warnings.filterwarnings("ignore")
-initialize(config_path="configs", version_base=None)
-
+import gradio as gr
+from hydra import compose, initialize
 from huggingface_hub import Repository
+from models.builder import build_model
+from segmentation.datasets import PascalVOCDataset
+from visualization import mask2rgb
 
-repo = Repository(
-    local_dir="clip-dinoiser",
-    clone_from="ariG23498/clip-dinoiser",
-    use_auth_token=os.environ.get("token")
-)
-
-check_path = 'clip-dinoiser/checkpoints/last.pt'
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-check = torch.load(check_path, map_location=device)
-dinoclip_cfg = "clip_dinoiser.yaml"
-cfg = compose(config_name=dinoclip_cfg)
-
-model = build_model(cfg.model, class_names=PascalVOCDataset.CLASSES).to(device)
-model.clip_backbone.decode_head.use_templates=False # switching off the imagenet templates for fast inference
-model.load_state_dict(check['model_state_dict'], strict=False)
-model = model.eval()
-
-import gradio as gr
+# Suppress warnings
+warnings.filterwarnings("ignore")
 
-colors = [
+# Constants
+CHECKPOINT_PATH = "clip-dinoiser/checkpoints/last.pt"
+CONFIG_PATH = "configs"
+DINOCLIP_CONFIG = "clip_dinoiser.yaml"
+COLORS = [
     (0, 255, 0),
-    (0, 0, 255),
-    (255, 255, 0),
-    (255, 0, 255),
+    (255, 0, 0),
     (0, 255, 255),
-    (114, 128, 250),
-    (0, 165, 255),
+    (255, 0, 255),
+    (255, 255, 0),
+    (250, 128, 114),
+    (255, 165, 0),
     (0, 128, 0),
     (144, 238, 144),
-    (238, 238, 175),
-    (255, 191, 0),
+    (175, 238, 238),
+    (0, 191, 255),
     (0, 128, 0),
-    (226, 43, 138),
+    (138, 43, 226),
     (255, 0, 255),
-    (0, 215, 255),
-    (255, 0, 0),
+    (255, 215, 0),
+    (0, 0, 255),
 ]
 
-color_map = {
-    f"{color_id}": f"#{hex(color[0])[2:].zfill(2)}{hex(color[1])[2:].zfill(2)}{hex(color[2])[2:].zfill(2)}" for color_id, color in enumerate(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")
+    )
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    checkpoint = torch.load(CHECKPOINT_PATH, map_location=device)
+    cfg = compose(config_name=DINOCLIP_CONFIG)
 
-def run_clip_dinoiser(input_image, text_prompts):
+    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)
+    return model.eval()
+
+def run_clip_dinoiser(input_image, text_prompts, model, device, colors):
     image = input_image.convert("RGB")
     text_prompts = text_prompts.split(",")
     palette = colors[:len(text_prompts)]
 
     model.clip_backbone.decode_head.update_vocab(text_prompts)
     model.to(device)
-    model.apply_found = True
 
     img_tens = T.PILToTensor()(image).unsqueeze(0).to(device) / 255.
-
     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)
 
+    mask = mask2rgb(output, palette)
     classes = np.unique(output).tolist()
-    palette_array = np.array(itemgetter(*classes)(palette)).reshape(1, -1, 3)
-    alpha=0.5
-    blend = (alpha)*np.array(image)/255. + (1-alpha) * mask/255.
+    alpha = 0.5
+    blend = (alpha * np.array(image) / 255.) + ((1 - alpha) * mask / 255.)
 
-    h_text = list()
-    for idx, text in enumerate(text_prompts):
-        h_text.append((text, f"{idx}"))
+    h_text = [(text, f"{idx}") for idx, text in enumerate(text_prompts)]
     return blend, mask, h_text
 
 
+def create_color_map(colors):
+    return {
+        f"{color_id}": f"#{hex(color[0])[2:].zfill(2)}{hex(color[1])[2:].zfill(2)}{hex(color[2])[2:].zfill(2)}"
+        for color_id, color in enumerate(colors)
+    }
 
 
-if __name__ == "__main__":
+def setup_gradio_interface(model, device, colors, color_map):
+    block = gr.Blocks()
 
-    block = gr.Blocks().queue()
     with block:
         gr.Markdown("<h1><center>CLIP-DINOiser<h1><center>")
 
@@ -106,15 +101,8 @@ if __name__ == "__main__":
                 run_button = gr.Button(value="Run")
 
             with gr.Column():
-                with gr.Row():
-                    overlay_mask = gr.Image(
-                        type="numpy",
-                        label="Overlay Mask",
-                    )
-                    only_mask = gr.Image(
-                        type="numpy",
-                        label="Segmentation Mask"
-                    )
+                overlay_mask = gr.Image(type="numpy", label="Overlay Mask")
+                only_mask = gr.Image(type="numpy", label="Segmentation Mask")
                 h_text = gr.HighlightedText(
                     label="Labels",
                     combine_adjacent=False,
@@ -123,16 +111,26 @@ if __name__ == "__main__":
                 )
 
         run_button.click(
-            fn=run_clip_dinoiser,
-            inputs=[input_image, text_prompts,],
+            fn=lambda img, prompts: run_clip_dinoiser(img, prompts, model, device, colors),
+            inputs=[input_image, text_prompts],
             outputs=[overlay_mask, only_mask, h_text]
         )
+
         gr.Examples(
-            [["vintage_bike.jpeg", "background, vintage bike, leather bag"]],
-            inputs = [input_image, text_prompts,],
-            outputs = [overlay_mask, only_mask, h_text],
-            fn=run_clip_dinoiser,
+            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),
             cache_examples=True,
             label='Try this example input!'
-      )
-    block.launch(share=False, show_api=False, show_error=True)
+        )
+
+    return block
+
+
+if __name__ == "__main__":
+    model = load_model()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    color_map = create_color_map(COLORS)
+    gradio_interface = setup_gradio_interface(model, device, COLORS, color_map)
+    gradio_interface.launch(share=False, show_api=False, show_error=True)

app_demo1.py

@@ -0,0 +1,138 @@
+# from models.builder import build_model
+# from visualization import mask2rgb
+# from segmentation.datasets import PascalVOCDataset
+
+# import os
+# from hydra import compose, initialize
+# from PIL import Image
+# import matplotlib.pyplot as plt
+# from torchvision import transforms as T
+# import torch.nn.functional as F
+# import numpy as np
+from operator import itemgetter
+# import torch
+# import random
+# import warnings
+
+warnings.filterwarnings("ignore")
+initialize(config_path="configs", version_base=None)
+
+# from huggingface_hub import Repository
+
+repo = Repository(
+    local_dir="clip-dinoiser",
+    clone_from="ariG23498/clip-dinoiser",
+    use_auth_token=os.environ.get("token")
+)
+
+check_path = 'clip-dinoiser/checkpoints/last.pt'
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+check = torch.load(check_path, map_location=device)
+dinoclip_cfg = "clip_dinoiser.yaml"
+cfg = compose(config_name=dinoclip_cfg)
+
+model = build_model(cfg.model, class_names=PascalVOCDataset.CLASSES).to(device)
+model.clip_backbone.decode_head.use_templates=False # switching off the imagenet templates for fast inference
+model.load_state_dict(check['model_state_dict'], strict=False)
+model = model.eval()
+
+# import gradio as gr
+
+colors = [
+    (0, 255, 0),
+    (0, 0, 255),
+    (255, 255, 0),
+    (255, 0, 255),
+    (0, 255, 255),
+    (114, 128, 250),
+    (0, 165, 255),
+    (0, 128, 0),
+    (144, 238, 144),
+    (238, 238, 175),
+    (255, 191, 0),
+    (0, 128, 0),
+    (226, 43, 138),
+    (255, 0, 255),
+    (0, 215, 255),
+    (255, 0, 0),
+]
+
+color_map = {
+    f"{color_id}": f"#{hex(color[0])[2:].zfill(2)}{hex(color[1])[2:].zfill(2)}{hex(color[2])[2:].zfill(2)}" for color_id, color in enumerate(colors)
+}
+
+def run_clip_dinoiser(input_image, text_prompts):
+    image = input_image.convert("RGB")
+    text_prompts = text_prompts.split(",")
+    palette = colors[:len(text_prompts)]
+
+    model.clip_backbone.decode_head.update_vocab(text_prompts)
+    model.to(device)
+    model.apply_found = True
+
+    img_tens = T.PILToTensor()(image).unsqueeze(0).to(device) / 255.
+
+    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 = output[0].argmax(dim=0)
+    mask = mask2rgb(output, palette)
+
+    classes = np.unique(output).tolist()
+    palette_array = np.array(itemgetter(*classes)(palette)).reshape(1, -1, 3)
+    alpha=0.5
+    blend = (alpha)*np.array(image)/255. + (1-alpha) * mask/255.
+
+    h_text = list()
+    for idx, text in enumerate(text_prompts):
+        h_text.append((text, f"{idx}"))
+    return blend, mask, h_text
+
+
+
+
+if __name__ == "__main__":
+
+    block = gr.Blocks().queue()
+    with block:
+        gr.Markdown("<h1><center>CLIP-DINOiser<h1><center>")
+
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(type="pil", label="Input Image")
+                text_prompts = gr.Textbox(label="Enter comma-separated prompts")
+                run_button = gr.Button(value="Run")
+
+            with gr.Column():
+                with gr.Row():
+                    overlay_mask = gr.Image(
+                        type="numpy",
+                        label="Overlay Mask",
+                    )
+                    only_mask = gr.Image(
+                        type="numpy",
+                        label="Segmentation Mask"
+                    )
+                h_text = gr.HighlightedText(
+                    label="Labels",
+                    combine_adjacent=False,
+                    show_legend=False,
+                    color_map=color_map
+                )
+
+        run_button.click(
+            fn=run_clip_dinoiser,
+            inputs=[input_image, text_prompts,],
+            outputs=[overlay_mask, only_mask, h_text]
+        )
+        gr.Examples(
+            [["vintage_bike.jpeg", "background, vintage bike, leather bag"]],
+            inputs = [input_image, text_prompts,],
+            outputs = [overlay_mask, only_mask, h_text],
+            fn=run_clip_dinoiser,
+            cache_examples=True,
+            label='Try this example input!'
+      )
+    block.launch(share=False, show_api=False, show_error=True)