SHOW/main.py

# -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright2023 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Authors: paper author. 
# Special Acknowlegement: Wojciech Zielonka and Justus Thies
# Contact: [email protected]

import mmcv
import argparse
import os
import SHOW
from stage1_main import SHOW_stage1
from stage2_main import SHOW_stage2
from configs.csv_parser import gen_path_from_ours_root
from SHOW import attr_dict
from pathlib import Path


def cvt_cfg(val):
 if val == 'True':
 return True
 elif val == 'False':
 return False
 elif val.isdigit():
 return int(val)
 else:
 return val


def parse_other_Cfg(other_cfg: list):
 parse_dict = {}
 parse_key = None
 parse_val = None

 def reg_key():
 nonlocal parse_dict
 nonlocal parse_key
 nonlocal parse_val
 if parse_key != None:
 parse_dict[parse_key] = parse_val
 parse_key = None
 parse_val = None

 for val in other_cfg:
 if val.startswith('--'):
 reg_key()
 parse_key = val.strip('--')
 else:
 if val.isdigit():
 val = int(val)
 if val == 'False':
 val = False
 if val == 'True':
 val = True
 if parse_val is None:
 parse_val = val
 elif isinstance(parse_val, list):
 parse_val.append(val)
 else:
 parse_val = [parse_val, val]
 reg_key()
 return parse_dict


def parse_overwrite_flag(over_write_cfg):

 over_write_cfg = over_write_cfg.strip(',')
 over_write_cfg = over_write_cfg.split(',')

 ret = {}
 for i in over_write_cfg:
 k, v = i.split('=')
 ret[k] = cvt_cfg(v)

 return ret


if __name__ == '__main__':
 SHOW.utils.work_seek_init()
 SHOW.utils.platform_init()

 parser = argparse.ArgumentParser()
 parser.add_argument('--speaker_name',
 type=str,
 default='-1',
 help='if speaker_name == -1, configs are auto loaded')
 parser.add_argument(
 '--all_top_dir',
 type=str,
 default='./test/half.mp4',
 )
 parser.add_argument('--others_cfg',
 type=str,
 nargs=argparse.REMAINDER,
 default=[])

 args = parser.parse_args()
 cfg = mmcv.Config(vars(args))

 cfg.over_write_cfg = parse_other_Cfg(cfg.others_cfg)
 cfg.all_top_dir = os.path.abspath(cfg.all_top_dir)

 tmp_dir = Path(cfg.all_top_dir)
 if tmp_dir.suffix.lower() in ['.mp4', '.avi', '.mov']:
 print(f'tmp_dir: {tmp_dir}')
 # all_top_dir/video.mp4
 parent = tmp_dir.parent
 cfg.all_top_dir = parent.__str__()
 img_path = parent.joinpath('image')

 if not os.path.exists(img_path):
 img_path.mkdir(parents=True, exist_ok=True)
 os.system(
 f'ffmpeg -i {tmp_dir} -vf fps=30 -q:v 1 {img_path}/%06d.png'
 )

 temp_cfg = gen_path_from_ours_root(speaker_name=cfg.speaker_name,
 all_top_dir=cfg.all_top_dir,
 **cfg.over_write_cfg)
 temp_cfg.update(cfg.over_write_cfg)
 temp_cfg = attr_dict(temp_cfg)
 temp_cfg.over_write_cfg = cfg.over_write_cfg

 if (temp_cfg.speaker_name == '-1' or temp_cfg.speaker_name == None):
 temp_cfg.speaker_name = -1

 SHOW_stage1(**temp_cfg)
 SHOW_stage2(**temp_cfg)