Create tools.py

utils/tools.py

@@ -0,0 +1,65 @@
+import cv2 
+import os 
+import sys 
+import clip 
+import numpy as np 
+from PIL import Image 
+import matplotlib.pyplot as plt 
+
+def convert_box_xywh_to_xyxy(box):
+    if len(box) == 4:
+        return [box[0],box[1],box[0]+box[2],box[1]+box[3]]
+    else:
+        result = []
+        for b in box:
+            b = convert_box_xywh_to_xyxy(b) 
+            result.append(b)
+    return result 
+
+
+def segment_image(image,bbox):
+    image_array = np.array(image) 
+    segmented_image_array = np.zeros_like(image_array)
+    x1,y1,x2,y2 = bbox 
+    segmented_image_array[y1:y2,x1:x2] = image_array[y1:y2,x1:x2]
+    segmented_image = Image.fromarray(segmented_image_array)
+    black_image = Image.new("RGB",image.size,(255,255,255))
+    transparency_mask = np.zeros((image_array.shape[0],image_array.shape[1]),dtype=np.uint8)
+    transparency_mask[y1:y2,x1:x2] = 255 
+    transparency_mask_image = Image.fromarray(transparency_mask,mode="L")
+    black_image.paste(segmented_image,mask=transparency_mask_image)
+    return black_image 
+
+def format_results(result,filter=0):
+    annotations = []
+    n = len(result.masks.data)
+    for i in range(n):
+        annotation = []
+        mask = result.masks.data[i] == 1.0 
+
+        if torch.sum(mask) < filter:
+            continue 
+        annotation['id'] = i 
+        annotation['segmentation'] = mask.cpu().numpy()
+        annotation['bbox'] = result.boxes.data[i]
+        annotation['score'] = result.boxes.conf[i]
+        annotation['area'] = annotation['segmentation'].sum()
+        annotations.append(annotation)
+    return annotations 
+
+def filter_masks(annotations):
+    annotations.sort(key=lambda x: x['area'],reverse=True)
+    to_remove = set() 
+    for i in range(0,len(annotations)):
+        a = annotations[i] 
+        for j in range(i+1,len(annotations)):
+            b = annotations[j] 
+            if i!=j and (j not in to_remove):
+                if b['area'] < a['area']:
+                    if (a['segmentation'] & b['segmentation']).sum()/b['segmentation'].sum()>0.8:
+                        to.remove.add(j)
+    return [a for i,a in enumerate(annotations) if i not in to_remove], to_remove 
+
+def get_bbox_from_mask(mask):
+    mask = mask.astype(np.uint8)
+    contours,hierarchy = cv2.findContours(mask,cv2.RETR)