Add script to convert alphapose-results to coco-in

This commit is contained in:
Ruben van de Ven 2023-01-28 21:41:15 +01:00
parent d575e64884
commit 5aa2d0fac1

View file

@ -0,0 +1,70 @@
"""
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
import logging
logging.basicConfig()
logger = logging.getLogger('loop_alphapose_training')
logger.setLevel(logging.INFO)
def coco_alphapose_merge_results(annotations_file: TextIOWrapper, results_file:TextIOWrapper, out_file: TextIOWrapper):
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())
results = json.loads(results_file.read())
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
json.dump(annotations,out_file)
logger.info(f'wrote to {out_file.name}')
if __name__ == "__main__":
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)