#!/usr/local/bin/python import argparse from common import * class Depth: def __init__(self, preset = None, architecture = None): self.preset = preset self.architecture = architecture def command(self, background_image, depth_image): AssertFileExists(background_image) command = "depth-estimation " + " --input=" + background_image + " --output=" + depth_image if self.architecture is not None: command = command + " --architecture=" + self.architecture if self.preset is not None: command = command + " --preset=" + self.preset return command class Slate: def __init__(self, architecture = None, slate_begin = None, slate_end = None): self.architecture = architecture self.slate_begin = slate_begin self.slate_end = slate_end def command(self, heartbeats, wide_video, wide_slate, tele_video = None, tele_slate = None): AssertFileExists(wide_video) command = "background-extraction " + " --input-video=" + wide_video + " --output-slate=" + wide_slate if self.architecture == "heartbeats" and FileExists(heartbeats): command = command + " --input-heartbeats=" + heartbeats if FileExists(tele_video) and (tele_slate is not None): command = command + " --input-video-multicam-tele=" + tele_video command = command + " --output-slate-multicam-tele=" + tele_slate return command class Matte: def __init__(self, architecture, model, num_frames = 0): self.architecture = architecture self.model = model self.num_frames = num_frames def command(self, input_video, output_video, background_image = None): command = "human-segmentation --input=" + input_video + " --output=" + output_video if FileExists(background_image): command = command + " --slate=" + background_image if self.architecture != None: command = command + " --architecture=" + self.architecture if self.model != None: command = command + " --model=" + self.model if self.num_frames is not None and self.num_frames > 0: command = command + " --num-frames=" + str(self.num_frames) return command class MatteAnalysis: def __init__(self): pass def command(self, input_video, input_matte, input_heartbeats, transparent_video=None): command = "python matte-analysis.py -iv {:s} -i {:s} -hb {:s}".format(input_video, input_matte, input_heartbeats) if transparent_video is not None: command += " -tv {:s}".format(transparent_video) return command ## ---- def GetCameraIdInDir(sequence_dir, content_name): wide_file_path = VideoFilePath(sequence_dir, content_name, "WA") tele_file_path = VideoFilePath(sequence_dir, content_name, "TP") if FileExists(wide_file_path) and FileExists(tele_file_path): return ["WA", "TP"] wide_file_path = VideoFilePath(sequence_dir, content_name, "1X") tele_file_path = VideoFilePath(sequence_dir, content_name, "2X") if FileExists(wide_file_path) and FileExists(tele_file_path): return ["1X", "2X"] for camera_id in ["WA", "1X", "TP", "2X", "TD", "LC", "FF", "UW"]: video_file_path = VideoFilePath(sequence_dir, content_name, camera_id) if FileExists( video_file_path ): return [camera_id] return [] def MatteFilePath(sequence_dir, sequence_name, camera_id, file_must_exist = True): segmentation_video = sequence_dir + sequence_name + "-" + camera_id + "-SM.MP4" if file_must_exist == False: return segmentation_video if FileExists(segmentation_video) and file_must_exist == True: return segmentation_video return "" def DepthFilePath(sequence_dir, sequence_name, camera_id, file_must_exist = True): background_image_depth = sequence_dir + sequence_name + "-" + camera_id + "-DM.PNG" if file_must_exist == False: return background_image_depth if FileExists(background_image_depth) and file_must_exist == True: return background_image_depth return "" def SlateFilePath(sequence_dir, sequence_name, camera_id, file_must_exist = True): background_image = sequence_dir + sequence_name + "-" + camera_id + "-BG.PNG" if file_must_exist == False: return background_image if FileExists(background_image) and file_must_exist == True: return background_image background_image = sequence_dir + sequence_name + "-" + camera_id + "-BG.JPG" if FileExists(background_image) and file_must_exist == True: return background_image return "" def VideoFilePath(sequence_dir, sequence_name, camera_id): video_file_path = sequence_dir + sequence_name + "-" + camera_id + ".MP4" if FileExists(video_file_path): return video_file_path video_file_path = sequence_dir + sequence_name + "-" + camera_id + ".MOV" if FileExists(video_file_path): return video_file_path return "" def HbeatFilePath(sequence_dir, sequence_name): input_heartbeats = sequence_dir + sequence_name + "-HB.ZIP" if FileExists(input_heartbeats): return input_heartbeats input_heartbeats = sequence_dir + sequence_name + "-HB.TXT" if FileExists(input_heartbeats): return input_heartbeats return "" ## ---- def sanitize_args(args): AssertDirExists(args.content) args.content = GetFullPath(args.content) if args.output_dir is not None: args.output_dir = GetFullPath(args.output_dir) return args def GetTestParams(content_dir, experiment_name): content_name = GetDirBaseName(content_dir) output_file_name = content_name if experiment_name != None: output_file_name = output_file_name + "-" + experiment_name if args.output_dir == "content": output_dir = content_dir elif args.output_dir is not None: output_dir = args.output_dir MakeDirRecursive(output_dir) return output_file_name, content_name, output_dir if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument('-c', '--content', type=str, default='', required=True, help='Content folder') parser.add_argument('-o', '--output-dir', type=str, default="./", required=False, help='Output (results) folder') parser.add_argument('-n', '--name', type=str, default=None, required=False, help='Name of experiment') parser.add_argument('-m', '--modules', type=str, default=None, nargs='+', required=True, choices=["matte", "slate", "depth"], help='Module to test') # matte params parser.add_argument('--matte-num-frames', type=int, default=None, help='Number of frames to process for matting') parser.add_argument('--matte-architecture', type=str, default=None, required=False, choices=["washington", "washington-base", "temporal"], help='Matte architecture to use.') parser.add_argument('--matte-model', type=str, default=None, required=False, choices=["resnet101", "resnet50", "mobilenetv2", "mobilenetv3"], help='Matte model') # slate params parser.add_argument('--slate-architecture', type=str, default="heartbeats", required=False, choices=["dlib", "heartbeats"], help='Slate architecture.') # depth params parser.add_argument('--depth-architecture', type=str, default=None, required=False, choices=["midas-v1-only", "multires-midas", "multires-structrl", "multires-leres"], help='Depth estimation architecture.') parser.add_argument('--depth-preset', type=str, default=None, required=False, choices=["precise", "fast", "veryfast"], help='Depth estimation architecture.') parser.add_argument('--log-timestamp', type=int, default=DefaultLogTimestamp(), help=LogTimestampHelp()) parser.add_argument('--log-level', type=str, default=DefaultLogLevel(), help=LogLevelHelp()) # matte analysis params parser.add_argument('--transparent-video', type=bool, default=None, help="Specify if transparent video be written to output") args = parser.parse_args() args = sanitize_args(args) InitLog(args.log_level, args.log_timestamp) # # printHeader(GetRealPath(__file__)) content_dirs = GetAllSubDirs(args.content) matte = Matte(args.matte_architecture, args.matte_model, args.matte_num_frames) slate = Slate(args.slate_architecture) depth = Depth(args.depth_preset, args.depth_architecture) # Slate command_list = [] if "slate" in args.modules: for content_dir in content_dirs: output_file_name, content_name, output_dir = GetTestParams(content_dir, args.name) camera_ids = GetCameraIdInDir(content_dir, content_name) if len(camera_ids) >= 1: hbeat_path = HbeatFilePath(content_dir, content_name) video_path = VideoFilePath(content_dir, content_name, camera_ids[0]) if FileExists(video_path) and FileExists(hbeat_path): command = "" if len(camera_ids) == 2: # Multicam case video_path_multicam_tele = VideoFilePath(content_dir, content_name, camera_ids[1]) if FileExists(video_path_multicam_tele): command = slate.command(hbeat_path, video_path, SlateFilePath(output_dir, output_file_name, camera_ids[0], False), video_path_multicam_tele, SlateFilePath(output_dir, output_file_name, camera_ids[1], False)) else: # Singlecam case command = slate.command(hbeat_path, video_path, SlateFilePath(output_dir, output_file_name, camera_ids[0], False)) if command != "": command_list.append({ 'command': command }) else: log.warning("Could not create command for " + content_dir) else: log.warning("Could not find any valid camera id in directory " + content_dir) ExecuteCommandList(command_list, GetNumberOfVirtualCores()) # Depth command_list = [] if "depth" in args.modules: for content_dir in content_dirs: output_file_name, content_name, output_dir = GetTestParams(content_dir, args.name) for camera_id in GetCameraIdInDir(content_dir, content_name): slate_path = SlateFilePath(content_dir, content_name, camera_id, True) if FileExists(slate_path): command = depth.command(slate_path, DepthFilePath(output_dir, output_file_name, camera_id, False)) command_list.append({ 'command': command }) ExecuteCommandList(command_list, 1) # Matte command_list = [] if "matte" in args.modules: for content_dir in content_dirs: output_file_name, content_name, output_dir = GetTestParams(content_dir, args.name) for camera_id in GetCameraIdInDir(content_dir, content_name): video_path = VideoFilePath(content_dir, content_name, camera_id) if FileExists(video_path): command = matte.command(video_path, MatteFilePath(output_dir, output_file_name, camera_id, False), SlateFilePath(content_dir, content_name, camera_id, True)) command_list.append({ 'command': command }) ExecuteCommandList(command_list, 1) # printFooter(GetRealPath(__file__), 0) # def TransparentVideoPath(sequence_dir, sequence_name, camera_id, file_must_exist = True): # transparent_video = sequence_dir + sequence_name + "-" + camera_id + "-MA.MP4" # if file_must_exist == False: # return transparent_video # if FileExists(transparent_video) and file_must_exist == True: # return transparent_video # return "" # TODO(Isam) - Why do you need matte analyzer in regression test? - Plz discuss with me # # Matte Analysis # command_list = [] # valid_camera_ids = ['WA', '1X'] # if "matte-analysis" in args.tests: # for content_dir in content_dirs: # output_file_name, content_name, output_dir = GetTestParams(content_dir, args.experiment) # for camera_id in GetAllCameraIds(): # if camera_id not in valid_camera_ids: # continue # video_file_path = VideoFilePath(content_dir, content_name, camera_id) # if args.transparent_video: # transparent_video_path = TransparentVideoPath(output_dir, output_file_name, camera_id, False) # if FileExists(video_file_path): # command = matte_analysis.command(video_file_path, MatteFilePath(content_dir, content_name, camera_id, True), \ # HbeatFilePath(content_dir, content_name), transparent_video=transparent_video_path) # command_list.append({'command': command, # 'stdout': LogFilePath(args.log, output_dir, output_file_name, camera_id, False)})