new detection format

scripts/detect_frames.py

@@ -77,10 +77,8 @@ def detect(in_seq_dir, out_seq_dir, model, device, verbose):
 
     with open(os.path.join(out_seq_dir, 'pred.json'), 'w') as f:
         json.dump(json_obj, f)
-    print('saved json')
-    
+
     torch.save(inference, 'inference.pt')
-    print('saved inference')
 
 def argument_parser():
     parser = argparse.ArgumentParser()

scripts/track_detection.py

@@ -1,3 +1,4 @@
+from lib.yolov5.utils.torch_utils import select_device
 import project_path
 from lib.yolov5.utils.general import clip_boxes, scale_boxes
 import argparse
@@ -66,21 +67,27 @@ def track_location(in_loc_dir, out_loc_dir, metadata_path, config, verbose):
             if (seq.startswith(".")): continue
             pbar.set_description("Processing " + seq)
 
-            in_path = os.path.join(in_loc_dir, seq, 'pred.json')
-            out_path = os.path.join(out_loc_dir, seq + ".txt")
 
-            track(in_path, out_path, metadata_path, seq, config, verbose)
+            track(in_loc_dir, out_loc_dir, metadata_path, seq, config, verbose)
 
-def track(in_path, out_path, metadata_path, seq_name, config, verbose):
+def track(in_loc_dir, out_loc_dir, metadata_path, seq_name, config, verbose):
+
+    json_path = os.path.join(in_loc_dir, seq, 'pred.json')
+    inference_path = os.path.join(in_loc_dir, seq, 'inference.pt')
+    out_path = os.path.join(out_loc_dir, seq + ".txt")
+
+
+    device_name = '0' if torch.cuda.is_available() else 'cpu'
+    device = select_device(device_name, batch_size=32)
+    inference = torch.load(inference_path, device=device)
+
+    print(type(inference))
+    print(len(inference))
+    print(type(inference[0]))
 
     # read detection
-    with open(in_path, 'r') as f:
+    with open(json_path, 'r') as f:
         detection = json.load(f)
-        print('inference' in detection)
-        print('image_shapes' in detection)
-        print(detection)
-
-    inference = detection['inference']
     image_shapes = detection['image_shapes']
     width = detection['width']
     height = detection['height']