human-segmentation/python_core/regression_tests.py

#!/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)})