import argparse import json import tqdm import cv2 import os import numpy as np from pycocotools import mask as mask_utils import random from PIL import Image from natsort import natsorted from pycocotools.mask import encode, decode, frPyObjects EVALMODE = "test" def blend_mask(input_img, binary_mask, alpha=0.5, color="g"): if input_img.ndim == 2: return input_img mask_image = np.zeros(input_img.shape, np.uint8) if color == "r": mask_image[:, :, 0] = 255 if color == "g": mask_image[:, :, 1] = 255 if color == "b": mask_image[:, :, 2] = 255 if color == "o": mask_image[:, :, 0] = 255 mask_image[:, :, 1] = 165 mask_image[:, :, 2] = 0 if color == "c": mask_image[:, :, 0] = 0 mask_image[:, :, 1] = 255 mask_image[:, :, 2] = 255 if color == "p": mask_image[:, :, 0] = 128 mask_image[:, :, 1] = 0 mask_image[:, :, 2] = 128 if color == "l": mask_image[:, :, 0] = 128 mask_image[:, :, 1] = 128 mask_image[:, :, 2] = 0 if color == "m": mask_image[:, :, 0] = 128 mask_image[:, :, 1] = 128 mask_image[:, :, 2] = 128 if color == "q": mask_image[:, :, 0] = 165 mask_image[:, :, 1] = 80 mask_image[:, :, 2] = 30 mask_image = mask_image * np.repeat(binary_mask[:, :, np.newaxis], 3, axis=2) blend_image = input_img[:, :, :].copy() pos_idx = binary_mask > 0 for ind in range(input_img.ndim): ch_img1 = input_img[:, :, ind] ch_img2 = mask_image[:, :, ind] ch_img3 = blend_image[:, :, ind] ch_img3[pos_idx] = alpha * ch_img1[pos_idx] + (1 - alpha) * ch_img2[pos_idx] blend_image[:, :, ind] = ch_img3 return blend_image def upsample_mask(mask, frame): H, W = frame.shape[:2] mH, mW = mask.shape[:2] if W > H: ratio = mW / W h = H * ratio diff = int((mH - h) // 2) if diff == 0: mask = mask else: mask = mask[diff:-diff] else: ratio = mH / H w = W * ratio diff = int((mW - w) // 2) if diff == 0: mask = mask else: mask = mask[:, diff:-diff] mask = cv2.resize(mask, (W, H)) return mask def downsample(mask, frame): H, W = frame.shape[:2] mH, mW = mask.shape[:2] mask = cv2.resize(mask, (W, H)) return mask #datapath /datasegswap #inference_path /inference_xmem_ego_last/coco #output /vis_piano #--show_gt要加上 if __name__ == "__main__": color = ['g', 'r', 'b', 'o', 'c', 'p', 'l', 'm', 'q'] data_path = "/data/work-gcp-europe-west4-a/yuqian_fu/datasets/HANDAL" output_path = "/data/work-gcp-europe-west4-a/yuqian_fu/datasets/HANDAL/Handal_vis_results_correct_last" mask_base_path = "/data/work-gcp-europe-west4-a/yuqian_fu/datasets/HANDAL/predictions_handal_all" json_path = "/data/work-gcp-europe-west4-a/yuqian_fu/datasets/HANDAL/handal_test_all_instruct_correct_videoname.json" with open(json_path, "r") as fp: datas = json.load(fp) print(len(datas)) video_select = ["handal_dataset_fixed_joint_pliers"] # debug for video_name in tqdm.tqdm(video_select): #实验需改动 data_list = [] for data in datas: if data["video_name"] == video_name: data_list.append(data) print(len(data_list)) data_list = random.sample(data_list, 100) for data in data_list: query_img = cv2.imread(os.path.join(data_path, data['first_frame_image'])) target_img = cv2.imread(os.path.join(data_path, data['image'])) # target_gt for i,ann in enumerate(data["anns"]): mask = decode(ann["segmentation"]) mask = downsample(mask, target_img) out = blend_mask(target_img, mask, color=color[0]) os.makedirs( f"{output_path}/{video_name}/target_gt", #debug exist_ok=True, ) img_path1 = data['image'] tmp_list = img_path1.split("/")[1:] joined_path = os.path.join(*tmp_list) #print("joined_path:", joined_path) output_file_path = os.path.join(output_path, video_name, "target_gt", joined_path) #print("output_file_path:", output_file_path) dir_path = output_file_path.split("/")[:-1] dir_path = "/".join(dir_path) os.makedirs( dir_path, #debug exist_ok=True,) cv2.imwrite( output_file_path, #debug out, ) # querty_gt for i,ann in enumerate(data["first_frame_anns"]): mask = decode(ann["segmentation"]) mask = downsample(mask, query_img) out = blend_mask(query_img, mask, color=color[0]) # os.makedirs( # f"{output_path}/{video_name}/query_gt", #debug # exist_ok=True, # ) new_path_query = data['first_frame_image'].replace(f"{video_name}/", "") output_file_path = os.path.join(output_path, video_name, "query_gt", new_path_query) #print("new_path_query:", new_path_query) #print("query_path:", f"{output_path}/{video_name}/query_gt/{new_path_query}") dir_path = output_file_path.split("/")[:-1] dir_path = "/".join(dir_path) os.makedirs( dir_path, #debug exist_ok=True,) cv2.imwrite( output_file_path, #debug out, ) # prediction mask_path = os.path.join(mask_base_path, data['image']) #print(mask_path) mask_path = mask_path.replace(".jpg", ".png") mask = Image.open(mask_path) mask = np.array(mask) unique_instances = np.unique(mask) unique_instances = unique_instances[unique_instances != 0] if len(unique_instances) > 9: continue for i,instance_value in enumerate(unique_instances): binary_mask = (mask == instance_value).astype(np.uint8) binary_mask = cv2.resize(binary_mask, (target_img.shape[1], target_img.shape[0])) try: binary_mask = upsample_mask(binary_mask, target_img) frame = blend_mask(target_img, binary_mask, color=color[i]) except: breakpoint() new_path_predict = data['image'].replace(f"{video_name}/", "") #print("new_path_predict:", new_path_predict) #print("predict_path:", f"{output_path}/{video_name}/predict/{new_path_predict}") output_file_path = os.path.join(output_path, video_name, "predict", new_path_predict) dir_path = output_file_path.split("/")[:-1] dir_path = "/".join(dir_path) os.makedirs( dir_path, #debug exist_ok=True,) cv2.imwrite( output_file_path, #debug frame, ) # os.makedirs( # f"{output_path}/{video_name}/predict", #debug # exist_ok=True, # ) # cv2.imwrite( # f"{output_path}/{video_name}/predict/{new_path_predict}", #debug # frame, # )