62 lines
1.7 KiB
Python
62 lines
1.7 KiB
Python
|
from pathlib import Path
|
||
|
|
import shutil
|
||
|
|
from loop_alphapose_training import run_cmd, Iteration
|
||
|
|
import logging
|
||
|
|
import coloredlogs
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
coloredlogs.install(level=logging.INFO)
|
||
|
|
|
||
|
|
def collate(iterations):
|
||
|
|
path = Path(f'out/test_imgs_{iterations[1].nr}/vis')
|
||
|
|
images = [x for x in path.iterdir() if x.is_file()]
|
||
|
|
for img in images:
|
||
|
|
logger.info(f"collate {img.name}")
|
||
|
|
target_dir = Path(f'out/test_imgs_loops/{img.stem}/')
|
||
|
|
target_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
for iteration in iterations:
|
||
|
|
try:
|
||
|
|
src = Path(f'out/test_imgs_{iteration.nr}/vis/') / img.name
|
||
|
|
target = target_dir / (iteration.nr + img.suffix)
|
||
|
|
logger.info(f"\tcopy {src} to {target}")
|
||
|
|
shutil.copy(src, target)
|
||
|
|
except Exception as e:
|
||
|
|
logger.exception(e)
|
||
|
|
|
||
|
|
|
||
|
|
# collate([Iteration(i) for i in range(2,7)])
|
||
|
|
# exit()
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
|
||
|
|
i = 0
|
||
|
|
iterations = []
|
||
|
|
while True:
|
||
|
|
i+=1
|
||
|
|
iteration = Iteration(i)
|
||
|
|
path = iteration.model_path()
|
||
|
|
if not path.exists():
|
||
|
|
logger.warning(f"Model for iteration {iteration.nr} doesn't exist")
|
||
|
|
break
|
||
|
|
iterations.append(iteration)
|
||
|
|
cmd = [
|
||
|
|
'python', 'scripts/demo_inference.py',
|
||
|
|
'--cfg', str(iteration.training_config_path()),
|
||
|
|
'--checkpoint', str(iteration.model_path()),
|
||
|
|
'--gpus', '0',
|
||
|
|
'--format', 'coco',
|
||
|
|
'--indir', 'data/test_imgs',
|
||
|
|
'--outdir', f'/out/test_imgs_{iteration.nr}',
|
||
|
|
'--save_img'
|
||
|
|
|
||
|
|
]
|
||
|
|
logger.info(f"Running {cmd}")
|
||
|
|
run_cmd(cmd, in_docker=True)
|
||
|
|
# break
|
||
|
|
|
||
|
|
|
||
|
|
collate(iterations)
|
||
|
|
|
||
|
|
|