File size: 2,264 Bytes
70c642e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 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) |