import math import cv2 import PIL import torch from PIL.Image import Image from typing import Union, Tuple, List, Optional import numpy as np import supervision as sv from sklearn.decomposition import PCA # def add_points_tag(img: Union[Image, np.ndarray], # point_labels: Union[List[int], np.ndarray] = None, # point_coords: Union[List[List[int]], np.ndarray] = None, # pil: bool = False): # if point_labels is None or point_coords is None or \ # not isinstance(point_labels, (List, np.ndarray)) or \ # not isinstance(point_coords, (List, np.ndarray)): # return img # if len(point_labels) != len(point_coords): # print('length of point_label and point_coordinate must be same!') # return img # if isinstance(img, Image): # img = np.uint8(img) # start_angle = 40 # x = 8 # y = 2 # def get_point(angle, d, base): # angle = angle / 180.0 * math.pi # _x, _y = math.cos(angle) * d, math.sin(angle) * d # return [base[0] + _x, base[1] - _y] # # assert len(point_labels) == len(point_coords), '' # for i in range(len(point_labels)): # points = [] # for j in range(5): # _x, _y = math.cos(start_angle), math.sin(start_angle) # points.append(get_point(start_angle, x, point_coords[i])) # start_angle -= 36 # points.append(get_point(start_angle, y, point_coords[i])) # start_angle -= 36 # points = np.array([points], np.int32) # color = (255, 0, 0) if point_labels[i] == 0 else (0, 255, 0) # cv2.fillPoly(img, points, color, cv2.LINE_AA) # if pil: # img = PIL.Image.fromarray(img) # return img def add_points_tag(img: Union[Image, np.ndarray], point_labels: Union[List[int], np.ndarray] = None, point_coords: Union[List[List[int]], np.ndarray] = None, pil: bool = False): if point_labels is None or point_coords is None or \ not isinstance(point_labels, (List, np.ndarray)) or \ not isinstance(point_coords, (List, np.ndarray)): return img if len(point_labels) != len(point_coords): print('length of point_label and point_coordinate must be same!') return img if isinstance(img, Image): img = np.array(img) # img.flags.writeable = True h, w = img.shape[:2] x_start_list, x_end_list = np.where((point_coords[:, 0] - 4) > 0, point_coords[:, 0] - 4, 0), np.where((point_coords[:, 0] + 4) < w, point_coords[:, 0] + 4, w) y_start_list, y_end_list = np.where((point_coords[:, 1] - 4) > 0, point_coords[:, 1] - 4, 0), np.where((point_coords[:, 1] + 4) < h, point_coords[:, 1] + 4, h) for i in range(len(point_labels)): x_start, x_end = x_start_list[i], x_end_list[i] y_start, y_end = y_start_list[i], y_end_list[i] label = point_labels[i] color = [0, 255, 0] if int(label) == 1 else [255, 0, 0] for x in range(x_start, x_end): for y in range(y_start, y_end): img[y, x, :] = color if pil: img = PIL.Image.fromarray(img) return img def add_boxes_tag(img: Union[Image, np.ndarray], boxes: Union[List[List[int]], np.ndarray] = None, pil: bool = False): if boxes is None or not isinstance(boxes, (List, np.ndarray)): return img # if isinstance(boxes, np.ndarray): # if not boxes.all(): # return img # else: # if not boxes: # return img if isinstance(img, Image): img = np.uint8(img) thickness = 2 for i in range(len(boxes)): color = (0, 255, 0) img = cv2.rectangle(img, (boxes[i][0], boxes[i][1]), (boxes[i][2], boxes[i][3]), color, thickness) if pil: img = PIL.Image.fromarray(img) return img def add_prompts_tag(img: Union[Image, np.ndarray], point_labels: Union[List[int], np.ndarray] = None, point_coords: Union[List[List[int]], np.ndarray] = None, boxes: Union[List[List[int]], np.ndarray] = None, pil: bool = False): img = add_points_tag(img, point_labels, point_coords, pil=pil) img = add_boxes_tag(img, boxes, pil=pil) return img def get_empty_detections(): detections = sv.Detections(xyxy=np.array([0, 0, 0, 0]).reshape(1, 4)) detections.xyxy = None return detections def pca_feature(feature: torch.Tensor, dim: int = 3, return_np: bool = True): pca = PCA(n_components=dim) H, W, C = feature.shape feature = feature.view(-1, C).cpu().numpy() feature = pca.fit_transform(feature) feature = torch.tensor(feature.reshape(H, W, dim)) if return_np: return feature.numpy() else: return feature def visual_feature_rgb(feature: torch.Tensor, pil:bool = True): assert feature.ndim >= 3, 'the dim of feature must >= 3!' if feature.ndim == 4: feature = feature.squeeze(0) if feature.shape[-1] != 3: feature = pca_feature(feature, 3, False) max_f, _ = feature.max(-1) min_f, _ = feature.min(-1) feature = (feature - min_f[..., None]) / (max_f[..., None] - min_f[..., None]) feature = np.uint8((feature*255).cpu().numpy()) if pil: return PIL.Image.fromarray(feature) else: return feature def transform_coords(src_shape, des_shape, points = None, boxes = None): assert points is not None or boxes is not None, 'one of points and boxes must be given!' scale_h = des_shape[0] / src_shape[0] scale_w = des_shape[1] / src_shape[1] if points is not None: new_points = np.full_like(points, 0) new_points[:, 0] = points[:, 0] * scale_w new_points[:, 1] = points[:, 1] * scale_h new_points.astype(np.int64) else: new_points = None if boxes is not None: new_boxes = np.full_like(boxes, 0) new_boxes[:, 0] = boxes[:, 0] * scale_w new_boxes[:, 1] = boxes[:, 1] * scale_h new_boxes[:, 2] = boxes[:, 2] * scale_w new_boxes[:, 3] = boxes[:, 3] * scale_h new_boxes.astype(np.int64) else: new_boxes = None return new_points, new_boxes def mask2greyimg(mask_list, pil=True): grey_img_list = [] for mask in mask_list: if pil: grey_img_list.append(PIL.Image.fromarray(np.uint8(mask*255))) else: grey_img_list.append(np.uint8(mask * 255)) return grey_img_list if __name__ == '__main__': src_shape = (100,100) des_shape = (200,200) points = np.array([[20,20],[40,40]]) boxes = np.array([[10,10,20,20]]) new_points, new_boxes = transform_coords(src_shape, des_shape, points, boxes) print(new_points, new_boxes)