chore: adding demo

app.py

@@ -11,6 +11,7 @@ import torch.nn.functional as F
 import numpy as np
 from operator import itemgetter 
 import torch
+import random
 import warnings
 
 warnings.filterwarnings("ignore")
@@ -19,9 +20,9 @@ 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")
+    local_dir="clip-dinoiser",
+    clone_from="ariG23498/clip-dinoiser",
+    use_auth_token=os.environ.get("token")
 )
 
 check_path = 'clip-dinoiser/checkpoints/last.pt'
@@ -38,8 +39,83 @@ model = model.eval()
 
 import gradio as gr
 
-def greet(name):
-    return "Hello " + name + "!!"
+def run_clip_dinoiser(input_image, text_prompts):
+    image = input_image.convert("RGB")
+    text_prompts = text_prompts.split(",")
+    palette = [
+        (random.randint(0, 256), random.randint(0, 256), random.randint(0, 256)) for _ in range(len(text_prompts))
+    ]
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+    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)
+
+    # fig = plt.figure(figsize=(3, 1))
+    # classes = np.unique(output).tolist()
+    # plt.imshow(np.array(itemgetter(*classes)(palette)).reshape(1, -1, 3))
+    # plt.xticks(np.arange(len(classes)), list(itemgetter(*classes)(text_prompts)), rotation=45)
+    # plt.yticks([])
+
+    
+    # fig, ax = plt.subplots(nrows=1, ncols=2)
+    # alpha=0.5
+    # blend = (alpha)*np.array(image)/255. + (1-alpha) * mask/255.
+    # ax[0].imshow(blend)
+    # ax[1].imshow(mask)
+    # ax[0].axis('off')
+    # ax[1].axis('off')
+
+    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.
+    return palette_array, blend, mask
+
+
+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(source='upload', type="pil")
+                text_prompts = gr.Textbox(label="Enter comma-separated prompts")
+                run_button = gr.Button(label="Run")
+
+            with gr.Column():
+                palette_array = gr.outputs.Image(
+                    type="numpy",
+                )
+                with gr.Row():
+                    overlay_mask = gr.outputs.Image(
+                        type="numpy",
+                    )
+                    only_mask = gr.outputs.Image(
+                        type="numpy",
+                    )
+
+        run_button.click(
+            fn=run_clip_dinoiser,
+            inputs=[input_image, text_prompts,],
+            outputs=[overlay_mask, only_mask]
+        )
+        gr.Examples(
+          [["vintage_bike.jpeg", "background, vintage bike, leather bag"]],
+          inputs = [input_image, text_prompts,],
+          outputs = [overlay_mask, only_mask],
+          fn=run_clip_dinoiser,
+          cache_examples=True,
+          label='Try this example input!'
+      )
+    block.launch(share=False, show_api=False, show_error=True)