Update detect_frames.py

scripts/detect_frames.py

@@ -68,20 +68,7 @@ def detect(in_seq_dir, out_seq_dir, model, device, verbose):
 
     inference, image_shapes, width, height = do_detection(dataloader, model, device, verbose=verbose)
 
-    print('converting tensor')
-    inference_list = []
-    with tqdm(total=len(inference), desc="...") as pbar:
-        for tensor in inference:
-            if isinstance(tensor, torch.Tensor):
-                inference_list.append(tensor.cpu().tolist())
-            else:
-                inference_list.append(tensor)
-            pbar.update(1)
-    pbar.close()
-
-    print('converted tensor')
     json_obj = {
-        'inference': inference_list,
         'image_shapes': image_shapes,
         'width': width,
         'height': height
@@ -90,7 +77,10 @@ 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')
+    print('saved json')
+    
+    torch.save(inference, 'inference.pt')
+    print('saved inference')
 
 def argument_parser():
     parser = argparse.ArgumentParser()