266 lines
9.7 KiB
Python
266 lines
9.7 KiB
Python
"""
|
|
TODO this script runs Alphapose's train.py, the created model is used to re-annotate the training-images, which is then fed back into the system
|
|
|
|
For now the only thing it does is that it merges alphapose-results.json with the coco input dataset.
|
|
"""
|
|
|
|
import argparse
|
|
import datetime
|
|
from io import TextIOWrapper
|
|
import json
|
|
from pathlib import Path
|
|
import subprocess
|
|
import yaml
|
|
import logging
|
|
|
|
logging.basicConfig()
|
|
logger = logging.getLogger('loop_alphapose_training')
|
|
logger.setLevel(logging.INFO)
|
|
|
|
class Iteration:
|
|
def __init__(self, i: int):
|
|
assert i > 0, "Iteration cannot be < 1"
|
|
self.i = i
|
|
self.nr = f"{i:04d}"
|
|
self.name = f"feedback_iteration{self.nr}"
|
|
|
|
def training_config_path(self):
|
|
return Path("./data/coco") / "_iterations" / self.nr / "256x192_res50_lr1e-3_1x.yaml"
|
|
|
|
def model_path(self):
|
|
assert self.i > 0, "Iteration 0 only used at training time"
|
|
return Path(f'exp/{self.name}-256x192_res50_lr1e-3_1x.yaml/final_DPG.pth')
|
|
|
|
def merged_results_path(self, stage: str):
|
|
assert stage in ['val', 'train']
|
|
return Path(f"./data/coco/_iterations/{self.nr}/alphapose-results-{stage}2017.json")
|
|
|
|
def interference_results_dir(self, stage: str, for_docker=False):
|
|
assert stage in ['val', 'train']
|
|
relative = "." if not for_docker else ""
|
|
return Path(f"{relative}/out/_iterations/{self.nr}_{stage}2017/")
|
|
|
|
def prev_iteration(self):
|
|
return Iteration(self.i - 1)
|
|
|
|
def next_iteration(self):
|
|
return Iteration(self.i+1)
|
|
|
|
@classmethod
|
|
def from_str(cls, input):
|
|
raise NotImplemented
|
|
|
|
def wrap_docker_cmd(cmd: list, container: str = 'alphapose'):
|
|
pwd = Path(__file__).parent.absolute()
|
|
return [
|
|
'docker', 'run',
|
|
'--rm',
|
|
'--gpus', 'all',
|
|
'--shm-size=10g',
|
|
'-v', str(pwd) + '/exp:/build/AlphaPose/exp',
|
|
'-v', str(pwd) + '/data:/build/AlphaPose/data',
|
|
'-v', str(pwd) + '/out:/out',
|
|
'-v', str(pwd) + '/alphapose-docker/detector/yolox/data:/build/AlphaPose/detector/yolox/data',
|
|
'-v', str(pwd) + '/alphapose-docker/detector/yolo/data:/build/AlphaPose/detector/yolo/data',
|
|
'-v', str(pwd) + '/alphapose-docker/pretrained_models:/build/AlphaPose/pretrained_models',
|
|
container,
|
|
*cmd
|
|
]
|
|
|
|
def create_config(iteration: Iteration):
|
|
"""
|
|
build config e.g. configs/coco/resnet/256x192_res50_lr1e-3_1x.yaml
|
|
"""
|
|
base_config = Path("../AlphaPose/configs/coco/resnet/256x192_res50_lr1e-3_1x.yaml")
|
|
with base_config.open('r') as fp:
|
|
config = yaml.safe_load(fp)
|
|
|
|
# the first iteration simply copies config
|
|
if iteration.i > 1:
|
|
prev_iteration = iteration.prev_iteration()
|
|
config['DATASET']['TRAIN']['ANN'] = str(prev_iteration.merged_results_path('train').relative_to(config['DATASET']['TRAIN']['ROOT']))
|
|
config['DATASET']['VAL']['ANN'] = str(prev_iteration.merged_results_path('val').relative_to(config['DATASET']['TRAIN']['ROOT']))
|
|
config['DATASET']['TEST']['DET_FILE'] = f'./exp/json/{iteration.name}_test_det_yolo.json'
|
|
config['DATASET']['TEST']['ANN'] = config['DATASET']['VAL']['ANN']
|
|
|
|
new_config = iteration.training_config_path()
|
|
|
|
if not new_config.parent.exists():
|
|
logger.info(f"Make directory for config: {new_config.parent}")
|
|
new_config.parent.mkdir()
|
|
|
|
with new_config.open('w') as fp:
|
|
yaml.dump(config, fp)
|
|
|
|
return new_config
|
|
|
|
def run_cmd(cmd, in_docker):
|
|
if in_docker:
|
|
cmd = wrap_docker_cmd(cmd)
|
|
|
|
logger.info(f"Run {cmd=}")
|
|
proc = subprocess.Popen (cmd, shell=False)
|
|
proc.communicate()
|
|
|
|
def create_and_run_training(iteration: Iteration):
|
|
'''
|
|
Basically just runs
|
|
python scripts/train.py
|
|
--cfg exp/config_first.yaml
|
|
--exp-id coco_test2
|
|
'''
|
|
|
|
create_config(iteration)
|
|
cmd = [
|
|
'python', 'scripts/train.py',
|
|
'--cfg', str(iteration.training_config_path()),
|
|
'--exp-id', iteration.name
|
|
]
|
|
|
|
run_cmd(cmd, in_docker = True)
|
|
|
|
def run_inferences(iteration: Iteration):
|
|
'''
|
|
create new train & validation datasets by basically running:
|
|
python scripts/demo_inference.py
|
|
--cfg configs/coco/resnet/256x192_res50_lr1e-3_1x.yaml
|
|
--checkpoint exp/coco_test1-256x192_res50_lr1e-3_1x.yaml/final_DPG.pth
|
|
--gpus 0
|
|
--indir data/coco/train2017
|
|
--outdir /out/first_train2017/
|
|
--format coco
|
|
--eval
|
|
python scripts/demo_inference.py
|
|
--cfg configs/coco/resnet/256x192_res50_lr1e-3_1x.yaml
|
|
--checkpoint exp/coco_test1-256x192_res50_lr1e-3_1x.yaml/final_DPG.pth
|
|
--gpus 0
|
|
--indir data/coco/val2017
|
|
--outdir /out/first_val2017/
|
|
--format coco
|
|
--eval
|
|
'''
|
|
|
|
base_cmd = [
|
|
'python', 'scripts/demo_inference.py',
|
|
'--cfg', str(iteration.training_config_path()),
|
|
'--checkpoint', str(iteration.model_path()),
|
|
'--gpus', '0',
|
|
'--format', 'coco',
|
|
'--eval',
|
|
]
|
|
|
|
cmd_infer_train = base_cmd.copy()
|
|
cmd_infer_train.extend([
|
|
'--indir', 'data/coco/train2017',
|
|
'--outdir', iteration.interference_results_dir('train', for_docker=True),
|
|
])
|
|
|
|
cmd_infer_val = base_cmd.copy()
|
|
cmd_infer_val.extend([
|
|
'--indir', 'data/coco/val2017',
|
|
'--outdir', iteration.interference_results_dir('val', for_docker=True),
|
|
])
|
|
|
|
run_cmd(cmd_infer_train, in_docker=True)
|
|
|
|
run_cmd(cmd_infer_val, in_docker=True)
|
|
|
|
def coco_alphapose_merge_results_for_iteration(iteration: Iteration):
|
|
for stage in ['val', 'train']:
|
|
coco_alphapose_merge_results(
|
|
Path(f"data/coco/annotations/person_keypoints_{stage}2017.json"),
|
|
iteration.interference_results_dir(stage) / "alphapose-results.json",
|
|
iteration.merged_results_path(stage)
|
|
)
|
|
|
|
def done_training(iteration: Iteration) -> bool:
|
|
return iteration.model_path().exists()
|
|
|
|
def done_interference(iteration: Iteration) -> bool:
|
|
return all([
|
|
(iteration.interference_results_dir(stage) / "alphapose-results.json").exists() for stage in ['train','val']
|
|
])
|
|
|
|
def done_merge_results(iteration: Iteration) -> bool:
|
|
return all([
|
|
iteration.merged_results_path(stage).exists() for stage in ['train','val']
|
|
])
|
|
|
|
def run_iteration(iteration: Iteration):
|
|
if not done_training(iteration):
|
|
create_and_run_training(iteration)
|
|
else:
|
|
logger.info(f"Training exists {iteration.name}")
|
|
if not done_interference(iteration):
|
|
run_inferences(iteration)
|
|
else:
|
|
logger.info(f"Inference results exist {iteration.name}")
|
|
if not done_merge_results(iteration):
|
|
coco_alphapose_merge_results_for_iteration(iteration)
|
|
else:
|
|
logger.info(f"Merged annotations exist {iteration.name}")
|
|
|
|
def coco_alphapose_merge_results(annotations_file: Path, results_file: Path, out_file: Path):
|
|
today = datetime.datetime.now().strftime("%Y/%m/%d")
|
|
info = {"description": "COCO 2017 Dataset, modified by Ruben van de Ven","url": "http://cocodataset.org","version": "0.1","year": 2023,"contributor": "COCO Consortium, Ruben van de Ven","date_created": today}
|
|
annotations = json.loads(annotations_file.read_text())
|
|
results = json.loads(results_file.read_text())
|
|
# annotations_ann:list = annotations['annotations']
|
|
|
|
# id_counts = {}
|
|
|
|
|
|
for i, result in enumerate(results):
|
|
if type(result['image_id']) == str:
|
|
result['image_id'] = int(result['image_id'][:-4])
|
|
|
|
result['id'] = i
|
|
result['iscrowd'] = 0 # TODO make sure this is a right terminology/assumption (what is this crowd here anyway individaul/crowd?)
|
|
result['bbox'] = result['box'] # TODO result.pop('box') to rename instead of copy
|
|
result['area'] = 1 # TODO : for now to bypass ignore in alphapose/datasets/mscoco.py:87
|
|
result['num_keypoints'] = 17 # TODO : verify that this is indeed always all points
|
|
|
|
# There can be multiple annotations per image. Try to match the originals by keeping track
|
|
# of their order of occurence
|
|
# if result['image_id'] not in id_counts:
|
|
# id_counts[result['image_id']] = 0
|
|
|
|
# # find matching annotations in original
|
|
# origs = list(filter(lambda ann: ann['image_id'] == result['image_id'], annotations_ann))
|
|
# assert len(origs) > id_counts[result['image_id']], f"Len should be one, found {len(origs)} for {result['image_id']}: {origs=}"
|
|
# orig = origs[id_counts[result['image_id']]]
|
|
# id_counts[result['image_id']] += 1
|
|
|
|
# result['id'] = orig['id'] # we keep track of the original id
|
|
|
|
annotations['annotations'] = results
|
|
annotations['info'] = info
|
|
|
|
with out_file.open('w') as fp:
|
|
json.dump(annotations, fp)
|
|
|
|
logger.info(f'wrote to {out_file.as_posix()}')
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
i = 1
|
|
while True:
|
|
iteration = Iteration(i)
|
|
logger.info(f"Run iteration {iteration.name}")
|
|
run_iteration(iteration)
|
|
i+=1
|
|
|
|
# parser = argparse.ArgumentParser(description='Merge alphapose-results.json with an input dataset')
|
|
# parser.add_argument('--annotations-file', required=True, type=argparse.FileType('r'),
|
|
# help='an annotations file from the COCO dataset (eg. person_keypoints_train2017.json)')
|
|
# parser.add_argument('--results-file', required=True, type=argparse.FileType('r'),
|
|
# help='path to the alphapose-results.json')
|
|
# parser.add_argument('--out-file', required=True, type=argparse.FileType('w'),
|
|
# help='the filename of the merged result')
|
|
|
|
|
|
# args = parser.parse_args()
|
|
|
|
# coco_alphapose_merge_results(args.annotations_file, args.results_file, args.out_file)
|