From f97a88870102e3b2a48d2881c11d478080ce7f37 Mon Sep 17 00:00:00 2001
From: sovit-123 <sovitrath5@gmail.com>
Date: Wed, 22 May 2024 19:18:54 +0530
Subject: [PATCH] Added JSON log to ONNX inference scripts

---
 inference.py            |  1 +
 onnx_inference_image.py | 11 +++++++++--
 onnx_inference_video.py | 11 +++++++++--
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/inference.py b/inference.py
index 61cdebd3..f760a226 100644
--- a/inference.py
+++ b/inference.py
@@ -298,6 +298,7 @@ def main(args):
     # Save JSON log file.
     if args['log_json']:
         log_json.save(os.path.join(OUT_DIR, 'log.json'))
+        
     # Calculate and print the average FPS.
     avg_fps = total_fps / frame_count
     print(f"Average FPS: {avg_fps:.3f}")
diff --git a/onnx_inference_image.py b/onnx_inference_image.py
index 4815a905..10392658 100644
--- a/onnx_inference_image.py
+++ b/onnx_inference_image.py
@@ -22,7 +22,7 @@
 from utils.annotations import (
     inference_annotations, convert_detections
 )
-# from utils.logging import log_to_json
+from utils.logging import LogJSON
 
 def collect_all_images(dir_test):
     """
@@ -136,6 +136,8 @@ def main(args):
     # score below this will be discarded.
     detection_threshold = args['threshold']
 
+    if args['log_json']:
+        log_json = LogJSON(os.path.join(OUT_DIR, 'log.json'))
 
     # To count the total number of frames iterated through.
     frame_count = 0
@@ -178,7 +180,7 @@ def main(args):
 
         # Log to JSON?
         if args['log_json']:
-            log_to_json(orig_image, os.path.join(OUT_DIR, 'log.json'), outputs)
+            log_json.update(orig_image, image_name, outputs[0], CLASSES)
 
         # Carry further only if there are detected boxes.
         if len(outputs[0]['boxes']) != 0:
@@ -208,6 +210,11 @@ def main(args):
 
     print('TEST PREDICTIONS COMPLETE')
     cv2.destroyAllWindows()
+
+    # Save JSON log file.
+    if args['log_json']:
+        log_json.save(os.path.join(OUT_DIR, 'log.json'))
+
     # Calculate and print the average FPS.
     avg_fps = total_fps / frame_count
     print(f"Average FPS: {avg_fps:.3f}")
diff --git a/onnx_inference_video.py b/onnx_inference_video.py
index 48c14dd1..82b0274c 100644
--- a/onnx_inference_video.py
+++ b/onnx_inference_video.py
@@ -26,7 +26,7 @@
 )
 from utils.transforms import infer_transforms, resize
 from deep_sort_realtime.deepsort_tracker import DeepSort
-# from utils.logging import log_to_json
+from utils.logging import LogJSON
 
 def read_return_video_data(video_path):
     cap = cv2.VideoCapture(video_path)
@@ -148,6 +148,9 @@ def main(args):
     else:
         RESIZE_TO = frame_width
 
+    if args['log_json']:
+        log_json = LogJSON(os.path.join(OUT_DIR, 'log.json'))
+
     frame_count = 0 # To count total frames.
     total_fps = 0 # To get the final frames per second.
 
@@ -185,7 +188,7 @@ def main(args):
 
             # Log to JSON?
             if args['log_json']:
-                log_to_json(frame, os.path.join(OUT_DIR, 'log.json'), outputs)
+                log_json.update(frame, save_name, outputs[0], CLASSES)
 
             # Carry further only if there are detected boxes.
             if len(outputs[0]['boxes']) != 0:
@@ -234,6 +237,10 @@ def main(args):
     # Close all frames and video windows.
     cv2.destroyAllWindows()
 
+    # Save JSON log file.
+    if args['log_json']:
+        log_json.save(os.path.join(OUT_DIR, 'log.json'))
+
     # Calculate and print the average FPS.
     avg_fps = total_fps / frame_count
     print(f"Average FPS: {avg_fps:.3f}")