import cv2 import matplotlib.pyplot as plt import numpy as np import torch from PIL import Image def segment_everything(_input, model, device, input_size=1024, iou_threshold=0.7, conf_threshold=0.25, better_quality=False, contour_thickness=1, max_det = 500): """ Performs segmentation on the input image and returns the segmented image and annotations. """ input_image = _input input_size = int(input_size) w, h = input_image.size scale = input_size / max(w, h) new_w = int(w * scale) new_h = int(h * scale) input_image = input_image.resize((new_w, new_h)) results = model(input_image, retina_masks=True, iou=iou_threshold, conf=conf_threshold, imgsz=input_size, max_det=max_det) annotations = results[0].masks.data segmented_image = fast_process(annotations=annotations, device=device, image=input_image, scale=(1024 // input_size), better_quality=better_quality, contour_thickness=contour_thickness) return segmented_image, annotations def fast_process(annotations, image, device, scale, better_quality=False, mask_random_color=True, bbox=None, use_retina=True, withContours=True, contour_thickness=2): if isinstance(annotations[0], dict): annotations = [annotation['segmentation'] for annotation in annotations] original_h = image.height original_w = image.width if better_quality: if isinstance(annotations[0], torch.Tensor): annotations = np.array(annotations.cpu()) for i, mask in enumerate(annotations): mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_OPEN, np.ones((3, 3), np.uint8)) annotations[i] = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_OPEN, np.ones((8, 8), np.uint8)) if device == 'cpu': annotations = np.array(annotations) inner_mask = fast_show_mask( annotations, plt.gca(), random_color=mask_random_color, bbox=bbox, retinamask=use_retina, target_height=original_h, target_width=original_w, ) else: if isinstance(annotations[0], np.ndarray): annotations = torch.from_numpy(annotations) inner_mask = fast_show_mask_gpu( annotations, plt.gca(), random_color=mask_random_color, bbox=bbox, retinamask=use_retina, target_height=original_h, target_width=original_w, ) if isinstance(annotations, torch.Tensor): annotations = annotations.cpu().numpy() kernel = np.ones((5, 5), np.uint8) if withContours: contour_all = [] temp = np.zeros((original_h, original_w, 1)) for i, mask in enumerate(annotations): if type(mask) == dict: mask = mask['segmentation'] annotation = mask.astype(np.uint8) # Perform morphological operations for separating connected objects and smoothing contours kernel = np.ones((5, 5), np.uint8) annotation = cv2.morphologyEx(annotation, cv2.MORPH_OPEN, kernel) annotation = cv2.GaussianBlur(annotation, (5, 5), 0) # Find contours contours, _ = cv2.findContours(annotation, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) for contour in contours: hull = cv2.convexHull(contour) epsilon = 0.005 * cv2.arcLength(contour, True) approx = cv2.approxPolyDP(contour, epsilon, True) contour_all.append(approx) for i, contour in enumerate(contour_all): M = cv2.moments(contour) if M["m00"] != 0: cX = int(M["m10"] / M["m00"]) cY = int(M["m01"] / M["m00"]) else: cX, cY = 0, 0 cv2.putText(temp, str(i), (cX, cY), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 125, 255), 2) cv2.drawContours(temp, contour_all, -1, (255, 255, 255), contour_thickness) color = np.array([255 / 255, 0 / 255, 0 / 255, 1]) # RGBA contour_mask = temp / 255 * color.reshape(1, 1, -1) image = image.convert('RGBA') overlay_inner = Image.fromarray((inner_mask * 255).astype(np.uint8), 'RGBA') image.paste(overlay_inner, (0, 0), overlay_inner) if withContours: overlay_contour = Image.fromarray((contour_mask * 255).astype(np.uint8), 'RGBA') image.paste(overlay_contour, (0, 0), overlay_contour) return image # CPU post process def fast_show_mask( annotation, ax, random_color=False, bbox=None, retinamask=True, target_height=960, target_width=960, ): mask_sum = annotation.shape[0] height = annotation.shape[1] weight = annotation.shape[2] # Sort annotation by area areas = np.sum(annotation, axis=(1, 2)) sorted_indices = np.argsort(areas)[::1] annotation = annotation[sorted_indices] index = (annotation != 0).argmax(axis=0) if random_color: color = np.random.random((mask_sum, 1, 1, 3)) else: color = np.ones((mask_sum, 1, 1, 3)) * np.array([30 / 255, 144 / 255, 255 / 255]) transparency = np.ones((mask_sum, 1, 1, 1)) * 0.6 visual = np.concatenate([color, transparency], axis=-1) mask_image = np.expand_dims(annotation, -1) * visual mask = np.zeros((height, weight, 4)) h_indices, w_indices = np.meshgrid(np.arange(height), np.arange(weight), indexing='ij') indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None)) mask[h_indices, w_indices, :] = mask_image[indices] if bbox is not None: x1, y1, x2, y2 = bbox ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor='b', linewidth=1)) if not retinamask: mask = cv2.resize(mask, (target_width, target_height), interpolation=cv2.INTER_NEAREST) return mask def fast_show_mask_gpu( annotation, ax, random_color=False, bbox=None, retinamask=True, target_height=960, target_width=960, ): device = annotation.device mask_sum = annotation.shape[0] height = annotation.shape[1] weight = annotation.shape[2] areas = torch.sum(annotation, dim=(1, 2)) sorted_indices = torch.argsort(areas, descending=False) annotation = annotation[sorted_indices] # Find the first non-zero value index for each position index = (annotation != 0).to(torch.long).argmax(dim=0) if random_color: color = torch.rand((mask_sum, 1, 1, 3)).to(device) else: color = torch.ones((mask_sum, 1, 1, 3)).to(device) * torch.tensor( [30 / 255, 144 / 255, 255 / 255] ).to(device) transparency = torch.ones((mask_sum, 1, 1, 1)).to(device) * 0.6 visual = torch.cat([color, transparency], dim=-1) mask_image = torch.unsqueeze(annotation, -1) * visual # Use vectorization to get the value of the batch mask = torch.zeros((height, weight, 4)).to(device) h_indices, w_indices = torch.meshgrid(torch.arange(height), torch.arange(weight)) indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None)) # Use vectorized indexing to update the show values mask[h_indices, w_indices, :] = mask_image[indices] mask_cpu = mask.cpu().numpy() if bbox is not None: x1, y1, x2, y2 = bbox ax.add_patch( plt.Rectangle( (x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="b", linewidth=1 ) ) if not retinamask: mask_cpu = cv2.resize( mask_cpu, (target_width, target_height), interpolation=cv2.INTER_NEAREST ) return mask_cpu