import torch from PIL import Image import cv2 import torchvision.transforms as T import os import numpy as np import torch.nn.functional as F from dift_util import DIFT_Demo, SDFeaturizer from torchvision.transforms import PILToTensor from tqdm import tqdm def load_dinov2(): dinov2_vitl14 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14').cuda() dinov2_vitl14.eval() return dinov2_vitl14 def infer_model(model, image): transform = T.Compose([ T.Resize((196, 196)), T.ToTensor(), T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) image = transform(image).unsqueeze(0).cuda() # cls_token = model.forward_features(image) cls_token = model(image, is_training=False) return cls_token def sort_frames(frame_name): return int(frame_name.split('.')[0]) 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) 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) # # 获取最大值即内接圆半径，中心点坐标 # minVal, maxVal, _, maxDistPt = cv2.minMaxLoc(raw_dist) # minVal = abs(minVal) # maxVal = abs(maxVal) return maxLoc, radius def pil_image_to_numpy(image, is_maks = False, index = 1): """Convert a PIL image to a NumPy array.""" if is_maks: image = image.resize((256, 256)) # 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((256, 256)) return np.array(image) def get_ID(images_list,masks_list,dinov2): 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') img_embedding = infer_model(dinov2, image) return img_embedding def get_dift_ID(feature_map,mask): # feature_map = feature_map * 0 new_feature = [] non_zero_coordinates = np.column_stack(np.where(mask != 0)) for coord in non_zero_coordinates: # feature_map[:, coord[0], coord[1]] = 1 new_feature.append(feature_map[:, coord[0], coord[1]]) stacked_tensor = torch.stack(new_feature, dim=0) # 在维度0上进行平均池化 average_pooled_tensor = torch.mean(stacked_tensor, dim=0) return average_pooled_tensor def extract_dift_feature(image, dift_model): if isinstance(image, Image.Image): image = image else: image = Image.open(image).convert('RGB') prompt = '' img_tensor = (PILToTensor()(image) / 255.0 - 0.5) * 2 dift_feature = dift_model.forward(img_tensor, prompt=prompt, up_ft_index=3,ensemble_size=8) return dift_feature dinov2 = load_dinov2() dinov2.requires_grad_(False) model_id = 'pretrained_models/chilloutmix' dift_model = SDFeaturizer(sd_id=model_id) # # 加载模型 # model = torch.load("/mmu-ocr/weijiawu/MovieDiffusion/ShowAnything/data/ref-youtube-vos/train/embedding/2cd01cf915/1.pth") # print(model.shape) # assert False dataset_type = "ref-youtube-vos" if dataset_type == "ref-youtube-vos": 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" save_p = "/mmu-ocr/weijiawu/MovieDiffusion/ShowAnything/data/ref-youtube-vos/train/embedding_SD_512_once" else: video_folder = "/mmu-ocr/weijiawu/MovieDiffusion/ShowAnything/data/VIPSeg/imgs" ann_folder = "/mmu-ocr/weijiawu/MovieDiffusion/ShowAnything/data/VIPSeg/panomasks" save_p = "/mmu-ocr/weijiawu/MovieDiffusion/ShowAnything/data/VIPSeg/embedding_SD_512_once" dataset_size = 512 dataset = [i for i in os.listdir(ann_folder)] for videoid in dataset: video_dir_1 = os.path.join(video_folder, videoid) ann_folder_1 = os.path.join(ann_folder, videoid) save_embedding = os.path.join(save_p, videoid) save_embedding_once = os.path.join(save_p, videoid+".pth") # if not os.path.exists(save_embedding): # print(save_embedding) # os.makedirs(save_embedding) image_files = sorted(os.listdir(video_dir_1), key=sort_frames) depth_files = sorted(os.listdir(ann_folder_1), key=sort_frames) #test keyframe_image = Image.open(os.path.join(video_dir_1, image_files[0])).convert('RGB') keyframe_image = keyframe_image.resize((dataset_size, dataset_size)) keyframe_dift = extract_dift_feature(keyframe_image, dift_model=dift_model) # torch.Size([1, 320, 32, 32]) mask = np.array(Image.open(os.path.join(ann_folder_1, depth_files[0]))) # np.array(Image.open(os.path.join(ann_folder_1, df))) # mask = Image.open(os.path.join(ann_folder_1, depth_files[0])).convert('P') ids = [i for i in np.unique(mask)] numpy_depth_images = np.array([pil_image_to_numpy(Image.open(os.path.join(ann_folder_1, df)),True,ids) for df in depth_files]) ids_list = {} for index_mask, mask in tqdm(enumerate(numpy_depth_images)): ids_embedding = torch.ones((dataset_size, dataset_size, 320)) # 判断文件是否存在 # if os.path.exists(os.path.join(save_embedding, '{}.pth'.format(index_mask))) and index_mask!=0: # continue for index in ids: mask_array = (np.array(mask)==index)*1 try: center_coordinate,_ = find_largest_inner_rectangle_coordinates(mask_array) except: continue print("find_largest_inner_rectangle_coordinates error") circle_img = np.zeros((dataset_size, dataset_size), np.float32) circle_mask = cv2.circle(circle_img, (center_coordinate[0],center_coordinate[1]), 20, 1, -1) # ID embedding if index_mask == 0: # diffusion feature mask_32 = cv2.resize(mask_array.astype(np.uint8),(int(dataset_size/8),int(dataset_size/8))) if len(np.column_stack(np.where(mask_32 != 0)))==0: continue id_feature = get_dift_ID(keyframe_dift[0],mask_32) ids_list[index] = id_feature else: try: id_feature = ids_list[index] except: print("index error") continue # 获取非零像素的坐标 # non_zero_coordinates = np.column_stack(np.where(circle_mask != 0)) # for coord in non_zero_coordinates: # ids_embedding[coord[0], coord[1]] = id_feature torch.save(ids_list, save_embedding_once) # only extract the feature of the first frame break # ids_embedding = F.avg_pool3d(ids_embedding, kernel_size=(2, 1, 1), stride=(8, 1, 1)) ids_embedding = F.avg_pool1d(ids_embedding, kernel_size=2, stride=2) torch.save(ids_embedding, os.path.join(save_embedding, '{}.pth'.format(index_mask)))