# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. #处理逻辑： #总体来说是视频分割逻辑，输入是第一帧的图片和query mask，输出是源源不断地后续预测结果 #1.若是memory，则将输入的第一帧设置为ego和ego mask，预测帧即后面的帧设置为exo的即可 #2.若不是memroy，针对每一预测target帧，找到定义vp_image和vp_mask的地方 import argparse import os from collections import defaultdict from pycocotools.mask import encode, decode, frPyObjects import numpy as np import torch from PIL import Image from sam2.build_sam import build_sam2_video_predictor import json from natsort import natsorted import cv2 import utils #from sklearn.metrics import balanced_accuracy_score parser = argparse.ArgumentParser() parser.add_argument( "--base_video_dir", type=str, default="/scratch/yuqian_fu/data_imgs", # debug help="base directory containing the videos to run VOS prediction on", ) parser.add_argument( "--sam2_cfg", type=str, default="configs/sam2/sam2_hiera_b+.yaml", help="SAM 2 model configuration file", ) parser.add_argument( "--sam2_checkpoint", type=str, default="./checkpoints/sam2_hiera_base_plus.pt", help="path to the SAM 2 model checkpoint", ) parser.add_argument( "--video_list_file", type=str, default=None, help="text file containing the list of video names to run VOS prediction on", ) parser.add_argument( "--output_mask_dir", type=str, required=True, help="directory to save the output masks (as PNG files)", ) parser.add_argument( "--score_thresh", type=float, default=0.0, help="threshold for the output mask logits (default: 0.0)", ) parser.add_argument( "--use_all_masks", action="store_true", help="whether to use all available PNG files in input_mask_dir " "(default without this flag: just the first PNG file as input to the SAM 2 model; " "usually we don't need this flag, since semi-supervised VOS evaluation usually takes input from the first frame only)", ) parser.add_argument( "--per_obj_png_file", action="store_true", help="whether use separate per-object PNG files for input and output masks " "(default without this flag: all object masks are packed into a single PNG file on each frame following DAVIS format; " "note that the SA-V dataset stores each object mask as an individual PNG file and requires this flag)", ) parser.add_argument( "--apply_postprocessing", action="store_true", help="whether to apply postprocessing (e.g. hole-filling) to the output masks " "(we don't apply such post-processing in the SAM 2 model evaluation)", ) parser.add_argument( "--track_object_appearing_later_in_video", action="store_true", help="whether to track objects that appear later in the video (i.e. not on the first frame; " "some VOS datasets like LVOS or YouTube-VOS don't have all objects appearing in the first frame)", ) parser.add_argument("--exoego", action='store_true', help="Use exoego dataset") # debug parser.add_argument("--start_id", type=str, default="0", help="Take ID to start with") # debug args = parser.parse_args() # the PNG palette for DAVIS 2017 dataset DAVIS_PALETTE = b"\x00\x00\x00\x80\x00\x00\x00\x80\x00\x80\x80\x00\x00\x00\x80\x80\x00\x80\x00\x80\x80\x80\x80\x80@\x00\x00\xc0\x00\x00@\x80\x00\xc0\x80\x00@\x00\x80\xc0\x00\x80@\x80\x80\xc0\x80\x80\x00@\x00\x80@\x00\x00\xc0\x00\x80\xc0\x00\x00@\x80\x80@\x80\x00\xc0\x80\x80\xc0\x80@@\x00\xc0@\x00@\xc0\x00\xc0\xc0\x00@@\x80\xc0@\x80@\xc0\x80\xc0\xc0\x80\x00\x00@\x80\x00@\x00\x80@\x80\x80@\x00\x00\xc0\x80\x00\xc0\x00\x80\xc0\x80\x80\xc0@\x00@\xc0\x00@@\x80@\xc0\x80@@\x00\xc0\xc0\x00\xc0@\x80\xc0\xc0\x80\xc0\x00@@\x80@@\x00\xc0@\x80\xc0@\x00@\xc0\x80@\xc0\x00\xc0\xc0\x80\xc0\xc0@@@\xc0@@@\xc0@\xc0\xc0@@@\xc0\xc0@\xc0@\xc0\xc0\xc0\xc0\xc0 \x00\x00\xa0\x00\x00 \x80\x00\xa0\x80\x00 \x00\x80\xa0\x00\x80 \x80\x80\xa0\x80\x80`\x00\x00\xe0\x00\x00`\x80\x00\xe0\x80\x00`\x00\x80\xe0\x00\x80`\x80\x80\xe0\x80\x80 @\x00\xa0@\x00 \xc0\x00\xa0\xc0\x00 @\x80\xa0@\x80 \xc0\x80\xa0\xc0\x80`@\x00\xe0@\x00`\xc0\x00\xe0\xc0\x00`@\x80\xe0@\x80`\xc0\x80\xe0\xc0\x80 \x00@\xa0\x00@ \x80@\xa0\x80@ \x00\xc0\xa0\x00\xc0 \x80\xc0\xa0\x80\xc0`\x00@\xe0\x00@`\x80@\xe0\x80@`\x00\xc0\xe0\x00\xc0`\x80\xc0\xe0\x80\xc0 @@\xa0@@ \xc0@\xa0\xc0@ @\xc0\xa0@\xc0 \xc0\xc0\xa0\xc0\xc0`@@\xe0@@`\xc0@\xe0\xc0@`@\xc0\xe0@\xc0`\xc0\xc0\xe0\xc0\xc0\x00 \x00\x80 \x00\x00\xa0\x00\x80\xa0\x00\x00 \x80\x80 \x80\x00\xa0\x80\x80\xa0\x80@ \x00\xc0 \x00@\xa0\x00\xc0\xa0\x00@ \x80\xc0 \x80@\xa0\x80\xc0\xa0\x80\x00`\x00\x80`\x00\x00\xe0\x00\x80\xe0\x00\x00`\x80\x80`\x80\x00\xe0\x80\x80\xe0\x80@`\x00\xc0`\x00@\xe0\x00\xc0\xe0\x00@`\x80\xc0`\x80@\xe0\x80\xc0\xe0\x80\x00 @\x80 @\x00\xa0@\x80\xa0@\x00 \xc0\x80 \xc0\x00\xa0\xc0\x80\xa0\xc0@ @\xc0 @@\xa0@\xc0\xa0@@ \xc0\xc0 \xc0@\xa0\xc0\xc0\xa0\xc0\x00`@\x80`@\x00\xe0@\x80\xe0@\x00`\xc0\x80`\xc0\x00\xe0\xc0\x80\xe0\xc0@`@\xc0`@@\xe0@\xc0\xe0@@`\xc0\xc0`\xc0@\xe0\xc0\xc0\xe0\xc0 \x00\xa0 \x00 \xa0\x00\xa0\xa0\x00 \x80\xa0 \x80 \xa0\x80\xa0\xa0\x80` \x00\xe0 \x00`\xa0\x00\xe0\xa0\x00` \x80\xe0 \x80`\xa0\x80\xe0\xa0\x80 `\x00\xa0`\x00 \xe0\x00\xa0\xe0\x00 `\x80\xa0`\x80 \xe0\x80\xa0\xe0\x80``\x00\xe0`\x00`\xe0\x00\xe0\xe0\x00``\x80\xe0`\x80`\xe0\x80\xe0\xe0\x80 @\xa0 @ \xa0@\xa0\xa0@ \xc0\xa0 \xc0 \xa0\xc0\xa0\xa0\xc0` @\xe0 @`\xa0@\xe0\xa0@` \xc0\xe0 \xc0`\xa0\xc0\xe0\xa0\xc0 `@\xa0`@ \xe0@\xa0\xe0@ `\xc0\xa0`\xc0 \xe0\xc0\xa0\xe0\xc0``@\xe0`@`\xe0@\xe0\xe0@``\xc0\xe0`\xc0`\xe0\xc0\xe0\xe0\xc0" root_path = "/scratch/yuqian_fu/data_imgs" # debug if not args.exoego: json_path = "/scratch/yuqian_fu/egoexo_val_framelevel_newprompt_all_instruction.json" else: json_path = "/scratch/yuqian_fu/ExoQuery_val_newprompt_all_instruction.json" with open(json_path, 'r') as f: datas = json.load(f) def fuse_davis_mask(mask_list): fused_mask = np.zeros_like(mask_list[0]) for mask in mask_list: fused_mask[mask != 0] = 1 return fused_mask def load_ann_png(path): """Load a PNG file as a mask and its palette.""" mask = Image.open(path) palette = mask.getpalette() mask = np.array(mask).astype(np.uint8) return mask, palette def save_ann_png(path, mask, palette): """Save a mask as a PNG file with the given palette.""" assert mask.dtype == np.uint8 assert mask.ndim == 2 output_mask = Image.fromarray(mask) output_mask.putpalette(palette) output_mask.save(path) def get_per_obj_mask(mask): """Split a mask into per-object masks.""" object_ids = np.unique(mask) object_ids = object_ids[object_ids > 0].tolist() per_obj_mask = {object_id: (mask == object_id) for object_id in object_ids} return per_obj_mask def put_per_obj_mask(per_obj_mask, height, width): """Combine per-object masks into a single mask.""" mask = np.zeros((height, width), dtype=np.uint8) object_ids = sorted(per_obj_mask)[::-1] for object_id in object_ids: object_mask = per_obj_mask[object_id] object_mask = object_mask.reshape(height, width) mask[object_mask] = object_id return mask #看看怎么获取调色板；或者调色版是否有必要；或者参考eval_davis里的调色板 def load_masks_from_dir( input_mask_dir, video_name, frame_name, per_obj_png_file, allow_missing=False ): """Load masks from a directory as a dict of per-object masks.""" if not per_obj_png_file: input_mask_path = os.path.join(input_mask_dir, video_name, f"{frame_name}.png") if allow_missing and not os.path.exists(input_mask_path): return {}, None input_mask, input_palette = load_ann_png(input_mask_path) per_obj_input_mask = get_per_obj_mask(input_mask) else: per_obj_input_mask = {} input_palette = None # each object is a directory in "{object_id:%03d}" format for object_name in os.listdir(os.path.join(input_mask_dir, video_name)): object_id = int(object_name) input_mask_path = os.path.join( input_mask_dir, video_name, object_name, f"{frame_name}.png" ) if allow_missing and not os.path.exists(input_mask_path): continue input_mask, input_palette = load_ann_png(input_mask_path) per_obj_input_mask[object_id] = input_mask > 0 return per_obj_input_mask, input_palette #ours #frame_name:"1" "2"类似 #这里其实最简单的只需要取出来第一帧ego视角下的mask即可 def load_masks_from_json(video_name, frame_name, per_obj_png_file, allow_missing=False): video_dir = os.path.join(root_path, video_name) data_list = [] for data in datas: if data["video_name"] == video_name: data_list.append(data) # 获取合适的第一帧 for data in data_list: if data['image'].split("/")[-1] == args.start_id + ".jpg": first_image = data break print("first_data:", first_image["first_frame_image"]) per_obj_input_mask = {} for ann in first_image["first_frame_anns"]: mask = decode(ann["segmentation"]) object_id = int(ann["category_id"]) per_obj_input_mask[object_id] = mask return per_obj_input_mask def save_masks_to_dir( output_mask_dir, video_name, frame_name, per_obj_output_mask, height, width, per_obj_png_file, output_palette, ): """Save masks to a directory as PNG files.""" os.makedirs(os.path.join(output_mask_dir, video_name), exist_ok=True) if not per_obj_png_file: output_mask = put_per_obj_mask(per_obj_output_mask, height, width) output_mask_path = os.path.join( output_mask_dir, video_name, f"{frame_name}.png" ) save_ann_png(output_mask_path, output_mask, output_palette) else: for object_id, object_mask in per_obj_output_mask.items(): object_name = f"{object_id:03d}" os.makedirs( os.path.join(output_mask_dir, video_name, object_name), exist_ok=True, ) output_mask = object_mask.reshape(height, width).astype(np.uint8) output_mask_path = os.path.join( output_mask_dir, video_name, object_name, f"{frame_name}.png" ) save_ann_png(output_mask_path, output_mask, output_palette) #memory机制针对的是同一段视频，而且这个脚本针对的是singel video，看看怎么扩展到若干个video #写一个大循环，把这个函数套进去 #看看怎么在不保存mask的情况下计算指标 @torch.inference_mode() @torch.autocast(device_type="cuda", dtype=torch.bfloat16) def vos_inference( base_video_dir, predictor, output_mask_dir, video_name, score_thresh=0.0, use_all_masks=False, per_obj_png_file=False, ): """Run VOS inference on a single video with the given predictor.""" # load the video frames and initialize the inference state on this video video_dir = os.path.join(base_video_dir, video_name) cams = os.listdir(video_dir) # cams.remove("annotation.json") # remove annotation file if exists for cam in cams: if "aria" in cam: ego = cam else: exo = cam print("ego exo", ego, exo) # debug if args.exoego: video_dir = os.path.join(video_dir, exo) print("video_dir:", video_dir) # debug else: video_dir = os.path.join(video_dir, ego) print("video_dir:", video_dir) # debug frame_names = [ os.path.splitext(p)[0] for p in os.listdir(video_dir) if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"] ] frame_names.sort(key=lambda p: int(os.path.splitext(p)[0])) #ours #video_dir:/data/work-gcp-europe-west4-a/yuqian_fu/Ego/data_segswap/3528e260-6a6d-46d7-b97d-b6c029ec7304 # missing_takes = 0 #记录丢失的takes # video_dir = os.path.join(root_path, video_name) # data_list = [] # frame_names = [] #frame_names存储的是帧数索引 # for data in datas: # if data["video_name"] == video_name: # data_list.append(data) # for data in data_list: # name = data["image"].split("/")[-1] # id = name.split(".")[0] # frame_names.append(id) # if len(data_list) == 0: # missing_takes += 1 # return [],[],[],[], missing_takes # data_tmp = data_list[0] # exo = data_tmp["image"].split("/")[-2] # ego = data_tmp["first_frame_image"].split("/")[-2] # print("ego exo",ego,exo) #debug # gt_path = f"{root_path}/{video_name}/annotation.json" # with open(gt_path, 'r') as fp: # gt = json.load(fp) # objs = list(gt['masks'].keys()) # objs_both_have = [] # for obj in objs: # if ego in gt["masks"][obj].keys() and exo in gt["masks"][obj].keys(): # objs_both_have.append(obj) # obj_ref = objs_both_have[0] # for obj in objs_both_have: # if len(list(gt["masks"][obj_ref][ego].keys())) < len(list(gt["masks"][obj][ego].keys())): # obj_ref = obj # IoUs = [] # ShapeAcc = [] # ExistenceAcc = [] # LocationScores = [] # all_ref_keys = np.asarray( # natsorted(gt["masks"][obj_ref][ego]) # ).astype(np.int64) # first_anno_key = str(all_ref_keys[0]) # obj_list_ego = [] # for obj in objs_both_have: # if first_anno_key in gt["masks"][obj][ego].keys(): # obj_list_ego.append(obj) #这里的video_dir是每个takes的路径 inference_state = predictor.init_state( video_path=video_dir, async_loading_frames=False ) height = inference_state["video_height"] width = inference_state["video_width"] input_palette = None # fetch mask inputs from input_mask_dir (either only mask for the first frame, or all available masks) # 仅利用第一帧的mask if not use_all_masks: # use only the first video's ground-truth mask as the input mask input_frame_inds = [0] # add those input masks to SAM 2 inference state before propagation object_ids_set = None for input_frame_idx in input_frame_inds: try: per_obj_input_mask = load_masks_from_json( video_name=video_name, frame_name=frame_names[input_frame_idx], per_obj_png_file=per_obj_png_file, ) except FileNotFoundError as e: raise RuntimeError( f"In {video_name=}, failed to load input mask for frame {input_frame_idx=}. " "Please add the `--track_object_appearing_later_in_video` flag " "for VOS datasets that don't have all objects to track appearing " "in the first frame (such as LVOS or YouTube-VOS)." ) from e # get the list of object ids to track from the first input frame if object_ids_set is None: object_ids_set = set(per_obj_input_mask) for object_id, object_mask in per_obj_input_mask.items(): # check and make sure no new object ids appear only in later frames if object_id not in object_ids_set: raise RuntimeError( f"In {video_name=}, got a new {object_id=} appearing only in a " f"later {input_frame_idx=} (but not appearing in the first frame). " "Please add the `--track_object_appearing_later_in_video` flag " "for VOS datasets that don't have all objects to track appearing " "in the first frame (such as LVOS or YouTube-VOS)." ) predictor.add_new_mask( inference_state=inference_state, frame_idx=input_frame_idx, obj_id=object_id, mask=object_mask, ) # check and make sure we have at least one object to track if object_ids_set is None or len(object_ids_set) == 0: raise RuntimeError( f"In {video_name=}, got no object ids on {input_frame_inds=}. " "Please add the `--track_object_appearing_later_in_video` flag " "for VOS datasets that don't have all objects to track appearing " "in the first frame (such as LVOS or YouTube-VOS)." ) # run propagation throughout the video and collect the results in a dict os.makedirs(os.path.join(output_mask_dir, video_name), exist_ok=True) # output_palette = input_palette or DAVIS_PALETTE output_palette = DAVIS_PALETTE video_segments = {} # video_segments contains the per-frame segmentation results #debug: 这里开始处理这个takes下的每一帧 for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video( inference_state ): per_obj_output_mask = { out_obj_id: (out_mask_logits[i] > score_thresh).cpu().numpy() for i, out_obj_id in enumerate(out_obj_ids) } id = frame_names[out_frame_idx] # gt_mask_list = [] # obj_list_exo = [] # for obj in obj_list_ego: # if id in gt["masks"][obj][exo].keys(): # obj_list_exo.append(obj) # for obj in obj_list_exo: # gt_mask = gt["masks"][obj][exo][id] # gt_mask = decode(gt_mask) # gt_mask_list.append(gt_mask) # if len(gt_mask_list) == 0: # continue # fused_gt_mask = fuse_davis_mask(gt_mask_list) pred_mask_list = list(per_obj_output_mask.values()) #ours pred_mask_list = [np.squeeze(mask, axis=0) for mask in pred_mask_list] if len(pred_mask_list) == 0: continue fused_pred_mask = fuse_davis_mask(pred_mask_list) h,w = fused_pred_mask.shape # gt_mask = cv2.resize(fused_gt_mask, (w, h), interpolation=cv2.INTER_NEAREST) # iou, shape_acc = utils.eval_mask(gt_mask, fused_pred_mask) # ex_acc = utils.existence_accuracy(gt_mask, fused_pred_mask) # location_score = utils.location_score(gt_mask, fused_pred_mask, size=(h, w)) # IoUs.append(iou) # ShapeAcc.append(shape_acc) # ExistenceAcc.append(ex_acc) # LocationScores.append(location_score) video_segments[out_frame_idx] = per_obj_output_mask for out_frame_idx, per_obj_output_mask in video_segments.items(): save_masks_to_dir( output_mask_dir=output_mask_dir, video_name=video_name, frame_name=frame_names[out_frame_idx], per_obj_output_mask=per_obj_output_mask, height=height, width=width, per_obj_png_file=per_obj_png_file, output_palette=output_palette, ) # write the output masks as palette PNG files to output_mask_dir # for out_frame_idx, per_obj_output_mask in video_segments.items(): # save_masks_to_dir( # output_mask_dir=output_mask_dir, # video_name=video_name, # frame_name=frame_names[out_frame_idx], # per_obj_output_mask=per_obj_output_mask, # height=height, # width=width, # per_obj_png_file=per_obj_png_file, # output_palette=output_palette, # ) # return IoUs.tolist(), ShapeAcc.tolist(), ExistenceAcc.tolist(), LocationScores.tolist(), missing_takes # return missing_takes @torch.inference_mode() @torch.autocast(device_type="cuda", dtype=torch.bfloat16) def vos_separate_inference_per_object( predictor, base_video_dir, input_mask_dir, output_mask_dir, video_name, score_thresh=0.0, use_all_masks=False, per_obj_png_file=False, ): """ Run VOS inference on a single video with the given predictor. Unlike `vos_inference`, this function run inference separately for each object in a video, which could be applied to datasets like LVOS or YouTube-VOS that don't have all objects to track appearing in the first frame (i.e. some objects might appear only later in the video). """ # load the video frames and initialize the inference state on this video video_dir = os.path.join(base_video_dir, video_name) frame_names = [ os.path.splitext(p)[0] for p in os.listdir(video_dir) if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"] ] frame_names.sort(key=lambda p: int(os.path.splitext(p)[0])) inference_state = predictor.init_state( video_path=video_dir, async_loading_frames=False ) height = inference_state["video_height"] width = inference_state["video_width"] input_palette = None # collect all the object ids and their input masks inputs_per_object = defaultdict(dict) for idx, name in enumerate(frame_names): if per_obj_png_file or os.path.exists( os.path.join(input_mask_dir, video_name, f"{name}.png") ): per_obj_input_mask, input_palette = load_masks_from_dir( input_mask_dir=input_mask_dir, video_name=video_name, frame_name=frame_names[idx], per_obj_png_file=per_obj_png_file, allow_missing=True, ) for object_id, object_mask in per_obj_input_mask.items(): # skip empty masks if not np.any(object_mask): continue # if `use_all_masks=False`, we only use the first mask for each object if len(inputs_per_object[object_id]) > 0 and not use_all_masks: continue print(f"adding mask from frame {idx} as input for {object_id=}") inputs_per_object[object_id][idx] = object_mask # run inference separately for each object in the video object_ids = sorted(inputs_per_object) output_scores_per_object = defaultdict(dict) for object_id in object_ids: # add those input masks to SAM 2 inference state before propagation input_frame_inds = sorted(inputs_per_object[object_id]) predictor.reset_state(inference_state) for input_frame_idx in input_frame_inds: predictor.add_new_mask( inference_state=inference_state, frame_idx=input_frame_idx, obj_id=object_id, mask=inputs_per_object[object_id][input_frame_idx], ) # run propagation throughout the video and collect the results in a dict for out_frame_idx, _, out_mask_logits in predictor.propagate_in_video( inference_state, start_frame_idx=min(input_frame_inds), reverse=False, ): obj_scores = out_mask_logits.cpu().numpy() output_scores_per_object[object_id][out_frame_idx] = obj_scores # post-processing: consolidate the per-object scores into per-frame masks os.makedirs(os.path.join(output_mask_dir, video_name), exist_ok=True) output_palette = input_palette or DAVIS_PALETTE video_segments = {} # video_segments contains the per-frame segmentation results for frame_idx in range(len(frame_names)): scores = torch.full( size=(len(object_ids), 1, height, width), fill_value=-1024.0, dtype=torch.float32, ) for i, object_id in enumerate(object_ids): if frame_idx in output_scores_per_object[object_id]: scores[i] = torch.from_numpy( output_scores_per_object[object_id][frame_idx] ) if not per_obj_png_file: scores = predictor._apply_non_overlapping_constraints(scores) per_obj_output_mask = { object_id: (scores[i] > score_thresh).cpu().numpy() for i, object_id in enumerate(object_ids) } video_segments[frame_idx] = per_obj_output_mask # write the output masks as palette PNG files to output_mask_dir for frame_idx, per_obj_output_mask in video_segments.items(): save_masks_to_dir( output_mask_dir=output_mask_dir, video_name=video_name, frame_name=frame_names[frame_idx], per_obj_output_mask=per_obj_output_mask, height=height, width=width, per_obj_png_file=per_obj_png_file, output_palette=output_palette, ) def main(): # if we use per-object PNG files, they could possibly overlap in inputs and outputs hydra_overrides_extra = [ "++model.non_overlap_masks=" + ("false" if args.per_obj_png_file else "true") ] predictor = build_sam2_video_predictor( config_file=args.sam2_cfg, ckpt_path=args.sam2_checkpoint, apply_postprocessing=args.apply_postprocessing, hydra_overrides_extra=hydra_overrides_extra, ) if args.use_all_masks: print("using all available masks in input_mask_dir as input to the SAM 2 model") else: print( "using only the first frame's mask in input_mask_dir as input to the SAM 2 model" ) # if a video list file is provided, read the video names from the file # (otherwise, we use all subdirectories in base_video_dir) split_path = "/home/yuqian_fu/Projects/ego-exo4d-relation/correspondence/SegSwap/data/split.json" with open(split_path, "r") as fp: data_split = json.load(fp) # video_names = data_split["val"] video_names = ["b511dfed-58f4-4c91-bf0a-f8ce9d47aea9"] # debug print(f"running VOS prediction on {len(video_names)} videos:\n{video_names}") # missing_num = 0 # total_iou = [] # total_shape_acc = [] # total_existence_acc = [] # total_location_scores = [] for n_video, video_name in enumerate(video_names): print(f"\n{n_video + 1}/{len(video_names)} - running on {video_name}") if not args.track_object_appearing_later_in_video: vos_inference( predictor=predictor, base_video_dir=args.base_video_dir, output_mask_dir=args.output_mask_dir, video_name=video_name, score_thresh=args.score_thresh, use_all_masks=args.use_all_masks, per_obj_png_file=args.per_obj_png_file, ) else: vos_separate_inference_per_object( predictor=predictor, base_video_dir=args.base_video_dir, input_mask_dir=args.input_mask_dir, output_mask_dir=args.output_mask_dir, video_name=video_name, score_thresh=args.score_thresh, use_all_masks=args.use_all_masks, per_obj_png_file=args.per_obj_png_file, ) # total_iou += ious # total_shape_acc += shape_accs # total_existence_acc += existence_accs # total_location_scores += location_scores # missing_num += missing_takes # print('TOTAL IOU: ', np.mean(total_iou)) # print('TOTAL LOCATION SCORE: ', np.mean(total_location_scores)) # print('TOTAL SHAPE ACC: ', np.mean(total_shape_acc)) # print("MISSING TAKES:", missing_num) print( f"completed VOS prediction on {len(video_names)} videos -- " f"output masks saved to {args.output_mask_dir}" ) if __name__ == "__main__": main()