| """Run the trained foreground model and UV parser with their existing environments.""" |
| import argparse |
| import json |
| import subprocess |
| import tempfile |
| import sys |
| from pathlib import Path |
|
|
|
|
| def main(): |
| root = Path(__file__).resolve().parents[1] |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument('--checkpoint', type=Path, required=True) |
| p.add_argument('--foreground-checkpoint', type=Path, required=True) |
| p.add_argument('--foreground-model-dir', type=Path, required=True) |
| p.add_argument('--foreground-python', default=sys.executable) |
| p.add_argument('--parser-python', default=sys.executable) |
| p.add_argument('--output-dir', type=Path, required=True) |
| p.add_argument('--probability-dir', type=Path, required=True) |
| p.add_argument('--pipeline', type=Path) |
| p.add_argument('--inputs', type=Path, nargs='+') |
| p.add_argument('--save-tensors', action='store_true') |
| opt = p.parse_args() |
| inputs = opt.inputs or [Path(x) for x in json.loads((root/'dense_uv_parser/v101_cases.json').read_text())] |
| for path in inputs + [opt.checkpoint, opt.foreground_checkpoint]: |
| if not path.is_file(): |
| raise FileNotFoundError(path) |
| entry = root/'dense_uv_parser/run_local.py' |
| with tempfile.TemporaryDirectory(prefix='foreground-batch-') as tmp: |
| manifest = Path(tmp)/'results.json' |
| subprocess.run([opt.foreground_python, str(entry), 'foreground_provider', |
| '--checkpoint', str(opt.foreground_checkpoint.resolve()), |
| '--model-dir', str(opt.foreground_model_dir.resolve()), |
| '--output-dir', str(opt.probability_dir.resolve()), |
| '--result-manifest', str(manifest), |
| '--inputs', *[str(x.resolve()) for x in inputs]], cwd=root, check=True) |
| probability = json.loads(manifest.read_text())['probabilities'] |
| if len(probability) != len(inputs): |
| raise ValueError('Foreground output count does not match inputs') |
| command = [opt.parser_python, str(entry), 'batch_accessories', |
| '--checkpoint', str(opt.checkpoint.resolve()), |
| '--output-dir', str(opt.output_dir.resolve()), |
| '--foreground-probabilities', *probability, |
| '--inputs', *[str(x.resolve()) for x in inputs]] |
| if opt.pipeline: |
| command += ['--pipeline', str(opt.pipeline.resolve())] |
| if opt.save_tensors: |
| command += ['--save-tensors'] |
| subprocess.run(command, cwd=root, check=True) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
