Updated inference

inference.py

@@ -24,8 +24,8 @@ WEIGHTS = 'models/v5m_896_300best.pt'
 # will need to configure these based on GPU hardware
 BATCH_SIZE = 32
 
-conf_thres = 0.001 # detection
-iou_thres  = 0.1 # NMS IOU
+conf_thres = 0.3 # detection
+iou_thres  = 0.3 # NMS IOU
 min_length = 0.3 # minimum fish length, in meters
 ###
 
@@ -72,7 +72,7 @@ def do_full_inference(dataloader, image_meter_width, image_meter_height, gp=None
         return
 
 
-    all_preds = do_suppression(dataloader, inference, width, height, gp=gp)
+    all_preds, real_width, real_height = do_suppression(dataloader, inference, width, height, gp=gp)
 
     results = do_tracking(all_preds, image_meter_width, image_meter_height, gp=gp)
 
@@ -122,7 +122,6 @@ def do_detection(dataloader, model, device, gp=None, batch_size=BATCH_SIZE):
             size = tuple(img.shape)
             nb, _, height, width = size  # batch size, channels, height, width
 
-            print(nb, _, height, width)
             # Run model & NMS
             with torch.no_grad():
                 inf_out, _ = model(img, augment=False) 
@@ -166,6 +165,8 @@ def do_suppression(dataloader, inference, width, height, gp=None, batch_size=BAT
                 # confidence score currently not used by tracker; set to 1.0
                 boxes = None
                 if box.shape[0]:
+                    real_width = shapes[si][0][1]
+                    real_height = shapes[si][0][0]
                     do_norm = partial(norm, w=shapes[si][0][1], h=shapes[si][0][0])
                     normed = list((map(do_norm, box[:, :4].tolist())))
                     boxes = np.stack([ [*bb, conf] for bb, conf in zip(normed, confs) ])
@@ -174,7 +175,7 @@ def do_suppression(dataloader, inference, width, height, gp=None, batch_size=BAT
 
             pbar.update(1*batch_size)
          
-    return all_preds
+    return all_preds, real_width, real_height
 
 def do_tracking(all_preds, image_meter_width, image_meter_height, gp=None): 
 

scripts/infer_frames.py

@@ -8,7 +8,6 @@ from aris import create_manual_marking, create_metadata_dictionary, prep_for_mm
 from inference import setup_model, do_suppression, do_detection, do_tracking, json_dump_round_float
 from visualizer import generate_video_batches
 import json
-from tracker import Tracker
 
 
 def main(args):
@@ -29,7 +28,7 @@ def main(args):
 
     dirname = args.frames
     
-    locations = ["kenai-val"]
+    locations = ["test"]
     for loc in locations:
 
         in_loc_dir = os.path.join(dirname, loc)
@@ -73,9 +72,15 @@ def infer_seq(in_dir, out_dir, seq_name, weights, metadata_path):
     # run detection + tracking
     model, device = setup_model(weights)
 
-    inference, width, height = do_detection(dataloader, model, device, gp=gradio_progress)
+    try:
+        inference, width, height = do_detection(dataloader, model, device, gp=gradio_progress)
+    except:
+        print("Error in " + seq_name)
+        with open(os.path.join(out_dir, "ERROR_" + seq_name + ".txt"), 'w') as f:
+            f.write("ERROR")
+        return
 
-    all_preds = do_suppression(dataloader, inference, width, height, gp=gradio_progress)
+    all_preds, real_width, real_height = do_suppression(dataloader, inference, width, height, gp=gradio_progress)
 
     results = do_tracking(all_preds, image_meter_width, image_meter_height, gp=gradio_progress)
 
@@ -84,12 +89,17 @@ def infer_seq(in_dir, out_dir, seq_name, weights, metadata_path):
         for fish in frame['fish']:
             bbox = fish['bbox']
             row = []
-            row.append(str(frame['frame_num']))
+            right = bbox[0]*real_width
+            top = bbox[1]*real_height
+            w = bbox[2]*real_width - bbox[0]*real_width
+            h = bbox[3]*real_height - bbox[1]*real_height
+
+            row.append(str(frame['frame_num'] + 1))
             row.append(str(fish['fish_id'] + 1))
-            row.append(str(int(bbox[0]*width)))
-            row.append(str(int(bbox[1]*height)))
-            row.append(str(int(bbox[2]*width)))
-            row.append(str(int(bbox[3]*height)))
+            row.append(str(int(right)))
+            row.append(str(int(top)))
+            row.append(str(int(w)))
+            row.append(str(int(h)))
             row.append("-1")
             row.append("-1")
             row.append("-1")