import os, io, csv, math, random import numpy as np from einops import rearrange import torch from decord import VideoReader import cv2 from scipy.ndimage import distance_transform_edt import torchvision.transforms as transforms from torch.utils.data.dataset import Dataset # from utils.util import zero_rank_print #from torchvision.io import read_image from PIL import Image import torchvision.transforms as T import torch.nn.functional as F def pil_image_to_numpy(image, is_maks = False, index = 1,size=256): """Convert a PIL image to a NumPy array.""" if is_maks: image = image.resize((size, size)) # image = (np.array(image)==index)*1 # image = cv2.cvtColor(image.astype(np.uint8), cv2.COLOR_GRAY2RGB) return np.array(image) else: if image.mode != 'RGB': image = image.convert('RGB') image = image.resize((size, size)) return np.array(image) def numpy_to_pt(images: np.ndarray, is_mask=False) -> torch.FloatTensor: """Convert a NumPy image to a PyTorch tensor.""" if images.ndim == 3: images = images[..., None] images = torch.from_numpy(images.transpose(0, 3, 1, 2)) if is_mask: return images.float() else: return images.float() / 255 def find_largest_inner_rectangle_coordinates(mask_gray): refine_dist = cv2.distanceTransform(mask_gray.astype(np.uint8), cv2.DIST_L2, 5, cv2.DIST_LABEL_PIXEL) _, maxVal, _, maxLoc = cv2.minMaxLoc(refine_dist) radius = int(maxVal) return maxLoc, radius # def find_largest_inner_rectangle_coordinates(mask_gray): # # 识别轮廓 # contours, _ = cv2.findContours(mask_gray.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) # xx,yy,ww,hh = 0,0,0,0 # contours_r = contours[0] # for contour in contours: # x, y, w, h = cv2.boundingRect(contour) # if w*h > ww*hh: # xx,yy,ww,hh = x, y, w, h # contours_r = contour # # 计算到轮廓的距离 # raw_dist = np.empty(mask_gray.shape, dtype=np.float32) # for i in range(mask_gray.shape[0]): # for j in range(mask_gray.shape[1]): # raw_dist[i, j] = cv2.pointPolygonTest(contours_r, (j, i), True) # # 获取最大值即内接圆半径，中心点坐标 # minVal, maxVal, _, maxDistPt = cv2.minMaxLoc(raw_dist) # minVal = abs(minVal) # maxVal = abs(maxVal) # return maxDistPt, int(maxVal) class YoutubeVos(Dataset): def __init__( self,video_folder,ann_folder,feature_folder, sample_size=512, sample_stride=4, sample_n_frames=14, ): self.dataset = [i.replace(".pth","") for i in os.listdir(feature_folder)] self.length = len(self.dataset) print(f"data scale: {self.length}") random.shuffle(self.dataset) self.video_folder = video_folder self.sample_stride = sample_stride self.sample_n_frames = sample_n_frames self.ann_folder = ann_folder self.heatmap = self.gen_gaussian_heatmap() self.feature_folder=feature_folder self.sample_size = sample_size print("length",len(self.dataset)) sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size) print("sample size",sample_size) def center_crop(self,img): h, w = img.shape[-2:] # Assuming img shape is [C, H, W] or [B, C, H, W] min_dim = min(h, w) top = (h - min_dim) // 2 left = (w - min_dim) // 2 return img[..., top:top+min_dim, left:left+min_dim] def gen_gaussian_heatmap(self,imgSize=200): circle_img = np.zeros((imgSize, imgSize), np.float32) circle_mask = cv2.circle(circle_img, (imgSize//2, imgSize//2), imgSize//2, 1, -1) isotropicGrayscaleImage = np.zeros((imgSize, imgSize), np.float32) # Guass Map for i in range(imgSize): for j in range(imgSize): isotropicGrayscaleImage[i, j] = 1 / 2 / np.pi / (40 ** 2) * np.exp( -1 / 2 * ((i - imgSize / 2) ** 2 / (40 ** 2) + (j - imgSize / 2) ** 2 / (40 ** 2))) isotropicGrayscaleImage = isotropicGrayscaleImage * circle_mask isotropicGrayscaleImage = (isotropicGrayscaleImage / np.max(isotropicGrayscaleImage)).astype(np.float32) isotropicGrayscaleImage = (isotropicGrayscaleImage / np.max(isotropicGrayscaleImage)*255).astype(np.uint8) # isotropicGrayscaleImage = cv2.resize(isotropicGrayscaleImage, (40, 40)) return isotropicGrayscaleImage def calculate_center_coordinates(self, numpy_images, masks, ids, feature_images,side=20): center_coordinates = [] ids_embedding_list = [] ids_list = {} for index_mask, mask in enumerate(masks): new_img = np.zeros((self.sample_size, self.sample_size), np.float32) ids_embedding = torch.zeros((self.sample_size, self.sample_size, 320)) # print(index_mask) 1024 576 for index in ids: mask_array = (np.array(mask)==index)*1 mask_32 = cv2.resize(mask_array.astype(np.uint8),(int(self.sample_size/8),int(self.sample_size/8))) if len(np.column_stack(np.where(mask_32 != 0)))==0: continue try: feature_image = feature_images[index] except: # print(feature_images.keys()) # print("KeyError: {}".format(index)) continue # 找到最大距离的索引 try: center_coordinate,radius = find_largest_inner_rectangle_coordinates(mask_array) side = int(radius) except: print("find_largest_inner_rectangle_coordinates error") continue x1 = max(center_coordinate[0]-side,1) x2 = min(center_coordinate[0]+side,self.sample_size-1) y1 = max(center_coordinate[1]-side,1) y2 = min(center_coordinate[1]+side,self.sample_size-1) if x2-x1<5 or y2-y1<5: continue need_map = cv2.resize(self.heatmap, (x2-x1, y2-y1)) new_img[y1:y2,x1:x2] = need_map if side>300: print("radius is too large") continue circle_img = np.zeros((self.sample_size, self.sample_size), np.float32) # try: circle_mask = cv2.circle(circle_img, (max(center_coordinate[0],1),min(center_coordinate[1],self.sample_size-1)), side, 1, -1) # except: # print((max(center_coordinate[0],1),min(center_coordinate[1],self.sample_size-1)), side) # 获取非零像素的坐标 non_zero_coordinates = np.column_stack(np.where(circle_mask != 0)) for coord in non_zero_coordinates: ids_embedding[coord[0], coord[1]] = feature_image # ID embedding # if index_mask == 0: # ids_list[index] = self.get_ID(numpy_images,mask_array) # 使用平均池化在第三个维度上进行池化，将大小减半 ids_embedding = F.avg_pool1d(ids_embedding, kernel_size=2, stride=2) new_img = cv2.cvtColor(new_img.astype(np.uint8), cv2.COLOR_GRAY2RGB) center_coordinates.append(new_img) ids_embedding_list.append(ids_embedding) return center_coordinates,ids_embedding_list def get_ID(self,images_list,masks_list): ID_images = [] image = images_list[0] mask = masks_list # 使用 findContours 函数找到轮廓 try: contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) x, y, w, h = cv2.boundingRect(contours[0]) mask = cv2.cvtColor(mask.astype(np.uint8), cv2.COLOR_GRAY2RGB) image = image * mask image = image[y:y+h,x:x+w] except: pass print("cv2.findContours error") # image = cv2.resize(image, (196, 196)) image = Image.fromarray(image).convert('RGB') image = self.idtransform(image).unsqueeze(0).to(dtype=torch.float16) image.to(self.device) # cls_token = self.dinov2(image, is_training=False) print(cls_token.shape) assert False # for i,m in zip(images_list,masks_list): # # image = self.idtransform(Image.fromarray(image)) # # cv2.imwrite("./vis/test.jpg", image) # ID_images.append(image) return ID_images def get_batch(self, idx): def sort_frames(frame_name): return int(frame_name.split('.')[0]) while True: videoid = self.dataset[idx] # videoid = video_dict['videoid'] preprocessed_dir = os.path.join(self.video_folder, videoid) ann_folder = os.path.join(self.ann_folder, videoid) feature_folder_file = os.path.join(self.feature_folder, videoid+".pth") if not os.path.exists(ann_folder): idx = random.randint(0, len(self.dataset) - 1) print("os.path.exists({}), error".format(ann_folder)) continue if not os.path.exists(feature_folder_file): idx = random.randint(0, len(self.dataset) - 1) print("os.path.exists({}), error".format(feature_folder_file)) continue # Sort and limit the number of image and depth files to 14 image_files = sorted(os.listdir(preprocessed_dir), key=sort_frames)[:self.sample_n_frames] depth_files = sorted(os.listdir(ann_folder), key=sort_frames)[:self.sample_n_frames] # feature_file = sorted(os.listdir(feature_folder_file), key=sort_frames)[:self.sample_n_frames] # Load image frames numpy_images = np.array([pil_image_to_numpy(Image.open(os.path.join(preprocessed_dir, img)),size=self.sample_size) for img in image_files]) pixel_values = numpy_to_pt(numpy_images) # Load feature frames feature_images = torch.load(feature_folder_file, map_location='cpu') # feature_images = np.array([np.array(torch.load(os.path.join(feature_folder_file, img))) for img in feature_file]) # feature_images = torch.tensor(feature_images).permute(0, 3, 1, 2) # Load mask frames mask = Image.open(os.path.join(ann_folder, depth_files[0])) ids = [i for i in np.unique(np.array(mask))] if len(ids)==1: idx = random.randint(0, len(self.dataset) - 1) print("len(ids), error") continue numpy_depth_images = np.array([pil_image_to_numpy(Image.open(os.path.join(ann_folder, df)),True,ids,size=self.sample_size) for df in depth_files]) heatmap_pixel_values,ids_embedding_list = self.calculate_center_coordinates(numpy_images,numpy_depth_images,ids,feature_images) ids_embedding_list = np.array([np.array(i) for i in ids_embedding_list]) ids_embedding_list = torch.from_numpy(ids_embedding_list.transpose(0, 3, 1, 2)) heatmap_pixel_values = np.array(heatmap_pixel_values) mask_pixel_values = numpy_to_pt(numpy_depth_images,True) heatmap_pixel_values = numpy_to_pt(heatmap_pixel_values,True) # Load motion values motion_values = 180 return pixel_values, mask_pixel_values, motion_values, heatmap_pixel_values, ids_embedding_list def __len__(self): return self.length def coordinates_normalize(self,center_coordinates): first_point = center_coordinates[0] center_coordinates = [one-first_point for one in center_coordinates] return center_coordinates def normalize(self, images): """ Normalize an image array to [-1,1]. """ return 2.0 * images - 1.0 def normalize_sam(self, images): """ Normalize an image array to [-1,1]. """ return (images - torch.tensor([0.485, 0.456, 0.406]).unsqueeze(0).unsqueeze(-1).unsqueeze(-1))/torch.tensor([0.229, 0.224, 0.225]).unsqueeze(0).unsqueeze(-1).unsqueeze(-1) def __getitem__(self, idx): pixel_values, mask_pixel_values,motion_values,heatmap_pixel_values,feature_images = self.get_batch(idx) pixel_values = self.normalize(pixel_values) sample = dict(pixel_values=pixel_values, mask_pixel_values=mask_pixel_values, motion_values=motion_values,heatmap_pixel_values=heatmap_pixel_values,Id_Images=feature_images) return sample def load_dinov2(): dinov2_vitl14 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14').cuda() dinov2_vitl14.eval() # dinov2_vitl14.requires_grad_(False) return dinov2_vitl14 if __name__ == "__main__": # from util import save_videos_grid # torch.multiprocessing.set_start_method('spawn') dino = load_dinov2() dino.to(dtype=torch.float16) dataset = YoutubeVos( video_folder = "/mmu-ocr/weijiawu/MovieDiffusion/ShowAnything/data/ref-youtube-vos/train/JPEGImages", ann_folder = "/mmu-ocr/weijiawu/MovieDiffusion/ShowAnything/data/ref-youtube-vos/train/Annotations", feature_folder = "/mmu-ocr/weijiawu/MovieDiffusion/ShowAnything/data/ref-youtube-vos/train/embedding", sample_size=256, sample_stride=1, sample_n_frames=16 ) # import pdb # pdb.set_trace() inverse_process = transforms.Compose([ transforms.Normalize(mean=[-0.485/0.229, -0.456/0.224, -0.406/0.225], std=[1/0.229, 1/0.224, 1/0.225]), ]) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, num_workers=10,) for idx, batch in enumerate(dataloader): images = ((batch["pixel_values"][0].permute(0,2,3,1)+1)/2)*255 masks = batch["mask_pixel_values"][0].permute(0,2,3,1)*255 heatmaps = batch["heatmap_pixel_values"][0].permute(0,2,3,1) # Id_Images = ((batch["Id_Images"][0])*torch.tensor([0.229, 0.224, 0.225]).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)+torch.tensor([0.485, 0.456, 0.406]).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)).permute(0,2,3,1)*255 # center_coordinates = batch["center_coordinates"] print(batch["pixel_values"].shape) # print(Id_Images.shape) for i in range(images.shape[0]): image = images[i].numpy().astype(np.uint8) # print(Id_Images[i].shape) # Id_Image = inverse_process(Id_Images[i]).permute(1,2,0).numpy().astype(np.uint8) # Id_Image = Id_Images[i].numpy().astype(np.uint8) # print(Id_Image.shape) mask = masks[i].numpy() heatmap = heatmaps[i].numpy() # center_coordinate = center_coordinates[i][0][:2].numpy().astype(np.uint8) # print(mask.shape) # print(center_coordinate) # mask[center_coordinate[0]:center_coordinate[0]+10,center_coordinate[1]:center_coordinate[1]+10]=125 print(np.unique(mask)) # print(Id_Image.shape) cv2.imwrite("./vis/image_{}.jpg".format(i), image) # cv2.imwrite("./vis/Id_Image_{}.jpg".format(i), Id_Image) cv2.imwrite("./vis/mask_{}.jpg".format(i), mask.astype(np.uint8)) cv2.imwrite("./vis/heatmap_{}.jpg".format(i), heatmap.astype(np.uint8)) cv2.imwrite("./vis/{}.jpg".format(i), heatmap.astype(np.uint8)*0.5+image*0.5) # save_videos_grid(batch["pixel_values"][i:i+1].permute(0,2,1,3,4), os.path.join(".", f"{idx}-{i}.mp4"), rescale=True) break