adding app with CLIP image segmentation

app.py

@@ -5,6 +5,8 @@ import numpy as np
 from PIL import Image
 import torch 
 import cv2 
+from matplotlib import pyplot as plt 
+from segmentation_mask_overlay import overlay_masks
 from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation,AutoProcessor,AutoConfig
 
 processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
@@ -33,17 +35,17 @@ def detect_using_clip(image,prompts=[],threshould=0.4):
     for i,prompt in enumerate(prompts):
         predicted_image =  torch.sigmoid(preds[i][0]).detach().cpu().numpy()
         predicted_image = np.where(predicted_image>threshould,255,0)
-        predicted_masks.append(create_rgb_mask(predicted_image))
-
-    return predicted_masks
+        predicted_masks.append(predicted_image)
+    bool_masks = [predicted_mask.astype('bool') for predicted_mask in predicted_masks]
+    return bool_masks
 
 def visualize_images(image,predicted_images,brightness=15,contrast=1.8):
     alpha = 0.7
     image_resize = cv2.resize(image,(352,352))
     resize_image_copy = image_resize.copy()
 
-    for mask_image in predicted_images:
-        resize_image_copy = cv2.addWeighted(resize_image_copy,alpha,mask_image,1-alpha,10)
+    # for mask_image in predicted_images:
+    #     resize_image_copy = cv2.addWeighted(resize_image_copy,alpha,mask_image,1-alpha,10)
 
     return cv2.convertScaleAbs(resize_image_copy, alpha=contrast, beta=brightness)     
 
@@ -52,10 +54,17 @@ def shot(brightness,contrast,image,labels_text):
         prompts = labels_text.split(',')
     else:
         prompts = [labels_text]
+    
     prompts = list(map(lambda x: x.strip(),prompts))
+
+    mask_labels = [f"{prompt}_{i}" for i,prompt in enumerate(prompts)]
+    cmap = plt.cm.tab20(np.arange(len(mask_labels)))[..., :-1]
+
+    resize_image = cv2.resize(image,(352,352))
+
     predicted_images  = detect_using_clip(image,prompts=prompts)
+    category_image = overlay_masks(resize_image,np.stack(predicted_images,-1),labels=mask_labels,colors=cmap,alpha=0.4,beta=1)
 
-    category_image = visualize_images(image=image,predicted_images=predicted_images,brightness=brightness,contrast=contrast)
     return category_image
 
 iface = gr.Interface(fn=shot,

requirements.txt

@@ -8,4 +8,6 @@ opencv-python
 Pillow
 requests
 urllib3<2
-git+https://github.com/facebookresearch/segment-anything.git
+git+https://github.com/facebookresearch/segment-anything.git
+segmentation_mask_overlay
+matplotlib