From 5aa2d0fac11d9af37c567284324d06b5a050d81d Mon Sep 17 00:00:00 2001
From: Ruben van de Ven <git@rubenvandeven.com>
Date: Sat, 28 Jan 2023 21:41:15 +0100
Subject: [PATCH] Add script to convert alphapose-results to coco-in

---
 loop_alphapose_training.py | 70 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 70 insertions(+)
 create mode 100644 loop_alphapose_training.py

diff --git a/loop_alphapose_training.py b/loop_alphapose_training.py
new file mode 100644
index 0000000..a107cdf
--- /dev/null
+++ b/loop_alphapose_training.py
@@ -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)
\ No newline at end of file