mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-17 00:23:14 +00:00
new inference script
This commit is contained in:
Executable
+111
@@ -0,0 +1,111 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from utils.config import Config
|
||||
from utils.registry_class import INFER_ENGINE
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Inference script",
|
||||
add_help=True
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image_file",
|
||||
type=str,
|
||||
default=None,
|
||||
required=True
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pose_directory",
|
||||
type=str,
|
||||
default=None,
|
||||
required=True
|
||||
)
|
||||
parser.add_argument(
|
||||
"--frame_directory",
|
||||
type=str,
|
||||
default=None,
|
||||
required=True
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pickle_data_file",
|
||||
type=str,
|
||||
default=None,
|
||||
required=True
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_frames",
|
||||
type=int,
|
||||
default=None,
|
||||
required=False
|
||||
)
|
||||
parser.add_argument(
|
||||
"--width",
|
||||
type=int,
|
||||
default=None,
|
||||
required=False
|
||||
)
|
||||
parser.add_argument(
|
||||
"--height",
|
||||
type=int,
|
||||
default=None,
|
||||
required=False
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=14,
|
||||
required=False
|
||||
)
|
||||
#parser.add_argument(
|
||||
# "--pretrained_model_name_or_path",
|
||||
# type=str,
|
||||
# default=None,
|
||||
# required=True
|
||||
#)
|
||||
args, _unknown = parser.parse_known_args()
|
||||
return args
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
cfg_update = Config(load=True)
|
||||
|
||||
# NB(bt,2025-02-04): Patching into their config-driven structure is super brittle. Be careful of updates.
|
||||
|
||||
if not args.image_file or not os.path.isfile(args.image_file):
|
||||
raise Exception(f"Image file does not exist: #{args.image_file}")
|
||||
|
||||
if not args.pose_directory or not os.path.isdir(args.pose_directory):
|
||||
raise Exception(f"Pose directory does not exist: #{args.pose_directory}")
|
||||
|
||||
if not args.frame_directory or not os.path.isdir(args.frame_directory):
|
||||
raise Exception(f"Frame directory does not exist: #{args.frame_directory}")
|
||||
|
||||
if not args.pickle_data_file or not os.path.isfile(args.pickle_data_file):
|
||||
raise Exception(f"Pickle data file does not exist: #{args.pickle_data_file}")
|
||||
|
||||
cfg_update.cfg_dict['test_list_path'] = [[
|
||||
4, # ???
|
||||
args.image_file,
|
||||
args.pose_directory,
|
||||
args.frame_directory,
|
||||
args.pickle_data_file,
|
||||
args.seed,
|
||||
]]
|
||||
|
||||
if args.max_frames:
|
||||
cfg_update.cfg_dict['max_frames'] = args.max_frames
|
||||
|
||||
if args.width and args.height:
|
||||
cfg_update.cfg_dict['resolution'] = [args.width, args.height]
|
||||
|
||||
print("Configurations:\n\n", cfg_update.cfg_dict, "\n\n")
|
||||
|
||||
INFER_ENGINE.build(dict(type=cfg_update.TASK_TYPE), cfg_update=cfg_update.cfg_dict)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -43,12 +43,17 @@ class Config(object):
|
||||
default=False,
|
||||
help='Into debug information'
|
||||
)
|
||||
parser.add_argument(
|
||||
"opts",
|
||||
help="other configurations",
|
||||
default=None,
|
||||
nargs=argparse.REMAINDER)
|
||||
return parser.parse_args()
|
||||
|
||||
# NB(bt,2025-02-04): We're parsing other arguments downstream of this.
|
||||
# parser.add_argument(
|
||||
# "opts",
|
||||
# help="other configurations",
|
||||
# default=None,
|
||||
# nargs=argparse.REMAINDER)
|
||||
# return parser.parse_args()
|
||||
args, _unknown = parser.parse_known_args()
|
||||
args.opts = []
|
||||
return args
|
||||
|
||||
def _path_join(self, path_list):
|
||||
path = ""
|
||||
|
||||
Reference in New Issue
Block a user