Update merge_eval_result.py

merge_eval_result.py

@@ -1,14 +1,47 @@
-import json, os
+import json, os, argparse
 
-eval_result_path = "file/example_eval_results"
-eval_result_files = [f for f in os.listdir(eval_result_path) if f.endswith('.json') and not 'merge' in f]
+def save_json(json_dict, fname):
+    with open(fname, "w") as f:
+        json.dump(json_dict, f, indent=4)
 
-merged_result = {}
-for fn in eval_result_files:
-    task_type = fn.replace('.json', '')
-    with open(f"{eval_result_path}/{fn}", "r") as f:
-        merged_result[task_type] = json.load(f)
+def load_json(fname):
+    with open(fname, "r") as f:
+        json_dict = json.load(f)
+    return json_dict
 
-merge_file = f"{eval_result_path}/merged_result.json"
-with open(merge_file, "w") as f:
-    json.dump(merged_result, f, indent=4)
+def merge(eval_result_path):
+    eval_result_files = [f for f in os.listdir(eval_result_path) if f.endswith('.json') and not 'merge' in f]
+    merged_result = {}
+    for fn in eval_result_files:
+        if 'tempcompass' in fn:
+            task_type = fn.replace('.json', '').replace('tempcompass_', '')
+            if task_type=='multi_choice':
+                task_type = 'multi-choice'
+            lmms_results = load_json(f"{eval_result_path}/{fn}")
+            results = {}
+            for lmms_result in lmms_results['logs']:
+                vid = lmms_result['doc']['video_id']
+                dim = lmms_result['doc']['dim']
+                if vid not in results:
+                    results[vid] = {}
+                if dim not in results[vid]:
+                    results[vid][dim] = []
+                
+                result = lmms_result['avg_accuracy']
+                result.pop('video_id')
+                result.pop('dim')
+                results[vid][dim].append(result)
+            merged_result[task_type] = results
+        else:
+            task_type = fn.replace('.json', '')
+            merged_result[task_type] = load_json(f"{eval_result_path}/{fn}")
+
+    merge_file = f"{eval_result_path}/merged_result.json"
+    save_json(merged_result, merge_file)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--eval_result_path", type=str, default="file/example_eval_results")
+    args = parser.parse_args()
+
+    merge(args.eval_result_path)