Inference config

inference.py

@@ -24,9 +24,12 @@ WEIGHTS = 'models/v5m_896_300best.pt'
 # will need to configure these based on GPU hardware
 BATCH_SIZE = 32
 
-conf_thres = 0.3 # detection
-iou_thres  = 0.3 # NMS IOU
-min_length = 0.3 # minimum fish length, in meters
+CONF_THRES = 0.3 # detection
+NMS_IOU  = 0.3 # NMS IOU
+MIN_LENGTH = 0.3 # minimum fish length, in meters
+MAX_AGE = 20 # time until missing fish get's new id
+IOU_THRES = 0.01 # IOU threshold for tracking
+MIN_HITS = 11 # minimum number of frames with a specific fish for it to count
 ###
 
 def norm(bbox, w, h):
@@ -131,7 +134,7 @@ def do_detection(dataloader, model, device, gp=None, batch_size=BATCH_SIZE):
 
     return inference, width, height
 
-def do_suppression(dataloader, inference, width, height, gp=None, batch_size=BATCH_SIZE):
+def do_suppression(dataloader, inference, width, height, gp=None, batch_size=BATCH_SIZE, conf_thres=CONF_THRES, iou_thres=NMS_IOU):
     """
     Args:
         frames_dir: a directory containing frames to be evaluated
@@ -177,7 +180,7 @@ def do_suppression(dataloader, inference, width, height, gp=None, batch_size=BAT
          
     return all_preds, real_width, real_height
 
-def do_tracking(all_preds, image_meter_width, image_meter_height, gp=None): 
+def do_tracking(all_preds, image_meter_width, image_meter_height, gp=None, max_age=MAX_AGE, iou_thres=IOU_THRES, min_hits=MIN_HITS, min_length=MIN_LENGTH): 
 
     if (gp): gp(0, "Tracking...")
 
@@ -188,7 +191,7 @@ def do_tracking(all_preds, image_meter_width, image_meter_height, gp=None):
         'image_meter_width': image_meter_width,
         'image_meter_height': image_meter_height
     }
-    tracker = Tracker(clip_info, args={ 'max_age': 20, 'min_hits': 0, 'iou_threshold': 0.01}, min_hits=11)
+    tracker = Tracker(clip_info, args={ 'max_age': max_age, 'min_hits': 0, 'iou_threshold': iou_thres}, min_hits=min_hits)
     
     # Run tracking
     with tqdm(total=len(all_preds), desc="Running tracking", ncols=0) as pbar:

scripts/inferEval.py

@@ -0,0 +1,34 @@
+import project_path
+import argparse
+from infer_frames import main as infer
+import sys
+sys.path.append('..')
+sys.path.append('../caltech-fish-counting')
+
+from evaluate import evaluate
+
+class Object(object):
+    pass
+
+def main(args):
+
+    infer_args = Object()
+    infer_args.metadata = "../caltech-fish-counting/data/metadata"
+    infer_args.frames = "../caltech-fish-counting/data/images"
+    infer_args.output = "../caltech-fish-counting/data/result"
+    infer_args.weights = "models/v5m_896_300best.pt"
+    infer_args.config = args.config
+
+    infer(infer_args)
+
+    evaluate("../frames/result_testing", "../frames/MOT", "../frames/metadata", "tracker", True)
+
+
+def argument_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", required=True, help="Config object. Required.")
+    return parser
+
+if __name__ == "__main__":
+    args = argument_parser().parse_args()
+    main(args)

scripts/infer_frames.py

@@ -26,9 +26,20 @@ def main(args):
     print("In task...")
     print("Cuda available in task?", torch.cuda.is_available())
 
+    # setup config
+    config = json.loads(args.config)
+    if "conf_threshold" not in config: config['conf_threshold'] = 0.3
+    if "nms_iou" not in config: config['nms_iou'] = 0.3
+    if "min_length" not in config: config['min_length'] = 0.3
+    if "max_age" not in config: config['max_age'] = 20
+    if "iou_threshold" not in config: config['iou_threshold'] = 0.01
+    if "min_hits" not in config: config['min_hits'] = 11
+
+    print(config)
+
     dirname = args.frames
     
-    locations = ["test"]
+    locations = ["kenai-val"]
     for loc in locations:
 
         in_loc_dir = os.path.join(dirname, loc)
@@ -39,6 +50,9 @@ def main(args):
         print(out_dir)
         print(metadata_path)
 
+        # run detection + tracking
+        model, device = setup_model(args.weights)
+
         seq_list = os.listdir(in_loc_dir)
         idx = 1
         for seq in seq_list:
@@ -47,11 +61,11 @@ def main(args):
             print(" ")
             idx += 1
             in_seq_dir = os.path.join(in_loc_dir, seq)
-            infer_seq(in_seq_dir, out_dir, seq, args.weights, metadata_path)
+            infer_seq(in_seq_dir, out_dir, config, seq, model, device, metadata_path)
 
-def infer_seq(in_dir, out_dir, seq_name, weights, metadata_path):
+def infer_seq(in_dir, out_dir, config, seq_name, model, device, metadata_path):
     
-    gradio_progress = lambda p, m: 0
+    #progress_log = lambda p, m: 0
 
     image_meter_width = -1
     image_meter_height = -1
@@ -68,21 +82,18 @@ def infer_seq(in_dir, out_dir, seq_name, weights, metadata_path):
 
     # create dataloader
     dataloader = create_dataloader_frames_only(in_dir)
-    
-    # run detection + tracking
-    model, device = setup_model(weights)
 
     try:
-        inference, width, height = do_detection(dataloader, model, device, gp=gradio_progress)
+        inference, width, height = do_detection(dataloader, model, device)
     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, real_width, real_height = do_suppression(dataloader, inference, width, height, gp=gradio_progress)
+    all_preds, real_width, real_height = do_suppression(dataloader, inference, width, height, conf_thres=config['conf_threshold'], iou_thres=config['nms_iou'])
 
-    results = do_tracking(all_preds, image_meter_width, image_meter_height, gp=gradio_progress)
+    results = do_tracking(all_preds, image_meter_width, image_meter_height, min_length=config['min_length'], max_age=config['max_age'], iou_thres=config['iou_threshold'], min_hits=config['min_hits'])
 
     mot_rows = []
     for frame in results['frames']:
@@ -118,6 +129,7 @@ def argument_parser():
     parser.add_argument("--frames", required=True, help="Path to frame directory. Required.")
     parser.add_argument("--metadata", required=True, help="Path to metadata directory. Required.")
     parser.add_argument("--output", required=True, help="Path to output directory. Required.")
+    parser.add_argument("--config", default="{}", help="Config object. Required.")
     parser.add_argument("--weights", default='models/v5m_896_300best.pt', help="Path to saved YOLOv5 weights. Default: ../models/v5m_896_300best.pt")
     return parser