modifided clip.py

extract_tools.py

@@ -187,7 +187,7 @@ def get_all_tools():
         clipseg_tool = Tool(
             name = 'ClipSegmentation-tool',
             func = clipsegmentation_mask,
-            description="""Use this tool when user ask to generate the segmentation Mask of the objects provided by the user. 
+            description="""Use this tool when user ask to extract the objects from the image . 
                         The input to the tool is the path of the image and list of objects for which Segmenation mask is to generated.
                         For example : 
                         Query :Provide a segmentation mask of all road car and dog in the image      
@@ -212,10 +212,10 @@ def get_all_tools():
         )
 
         object_extractor = Tool(
-            name = "Object Extraction Tool",
+            name = "Object description Tool",
             func = object_extraction,
-            description = " The Tool is used to extract objects within the image . Use this tool if user specifically ask to identify \
-                what are the objects I can view in the image or identify the objects within the image . "
+            description = " The Tool is used to describe the objects within the image . Use this tool if user specifically ask to identify \
+                what are the objects I can view in the image or identify the objects within the image. "
         )
 
         image_parameters_tool = Tool(

tool_utils/clip_segmentation.py

@@ -14,11 +14,22 @@ class CLIPSEG:
         self.threshould = threshould
         self.clip_model.to('cpu')
     
+    @ staticmethod 
+    def create_single_mask(predicted_masks , color = None ):
+        
+        if len(predicted_masks)>0:
+            mask_image = np.zeros_like(predicted_masks[0])
+        else:
+            mask_image = np.zeros(shape=(352,352),dtype=np.unit8)
+        for masks in predicted_masks:
+            mask_image = np.bitwise_or(mask_image,masks)
+        return mask_image
+
     @staticmethod
     def create_rgb_mask(mask,color=None):
-        color = tuple(np.random.choice(range(0,256), size=3))
+        color = tuple(np.random.choice(range(128,255), size=3))
         gray_3_channel = cv2.merge((mask, mask, mask))
-        gray_3_channel[mask==255] = color
+        gray_3_channel[mask==255] = 255 # for orignial color
         return gray_3_channel.astype(np.uint8)
     
     def get_segmentation_mask(self,image_path:str,object_prompts:List):
@@ -41,16 +52,25 @@ class CLIPSEG:
             predicted_mask =  torch.sigmoid(preds[i][0]).detach().cpu().numpy()
             predicted_mask = np.where(predicted_mask>self.threshould, 255,0)
             predicted_masks.append(predicted_mask)
-
+        
+        final_mask = self.create_single_mask(predicted_masks)
+        rgb_predicted_mask = self.create_rgb_mask(final_mask)
+        
         resize_image = cv2.resize(image,(352,352))
-        mask_labels = [f"{prompt}_{i}" for i,prompt in enumerate(object_prompts)]
-        cmap = plt.cm.tab20(np.arange(len(mask_labels)))[..., :-1]
+        rgb_mask_img = cv2.bitwise_and(resize_image,rgb_predicted_mask )
+        
+        # mask_labels = [f"{prompt}_{i}" for i,prompt in enumerate(object_prompts)]
+        # cmap = plt.cm.tab20(np.arange(len(mask_labels)))[..., :-1]
 
-        bool_masks = [predicted_mask.astype('bool') for predicted_mask in predicted_masks]
-        final_mask = overlay_masks(resize_image,np.stack(bool_masks,-1),labels=mask_labels,colors=cmap,alpha=0.5,beta=0.7)
+        # bool_masks = [predicted_mask.astype('bool') for predicted_mask in predicted_masks]
+        # final_mask = overlay_masks(resize_image,np.stack(bool_masks,-1),labels=mask_labels,colors=cmap,alpha=0.5,beta=0.7)
         try:
-            cv2.imwrite('final_mask.png',final_mask)
+            cv2.imwrite('final_mask.png',rgb_mask_img)
             return 'Segmentation image created : final_mask.png'
         except Exception as e:
             logging.error("Error while saving the final mask :",e)
-            return "unable to create a mask image "
+            return "unable to create a mask image "
+    
+if __name__=="__main__": 
+  clip = CLIPSEG()
+  obj = clip.get_segmentation_mask(image_path="../image_store/demo.jpg",object_prompts=['sand','dog'])