inference

inference.py

@@ -102,7 +102,7 @@ def setup_model(weights_fp=WEIGHTS, imgsz=896, batch_size=32):
     
     return model, device
                        
-def do_detection(dataloader, model, device, gp=None, batch_size=BATCH_SIZE):
+def do_detection(dataloader, model, device, gp=None, batch_size=BATCH_SIZE, verbose=True):
     """
     Args:
         frames_dir: a directory containing frames to be evaluated
@@ -115,7 +115,7 @@ def do_detection(dataloader, model, device, gp=None, batch_size=BATCH_SIZE):
 
     inference = []
     # Run detection
-    with tqdm(total=len(dataloader)*batch_size, desc="Running detection", ncols=0) as pbar:
+    with tqdm(total=len(dataloader)*batch_size, desc="Running detection", ncols=0, disable=not verbose) as pbar:
         for batch_i, (img, _, shapes) in enumerate(dataloader):
 
             if gp: gp(batch_i / len(dataloader), pbar.__str__())
@@ -134,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, conf_thres=CONF_THRES, iou_thres=NMS_IOU):
+def do_suppression(dataloader, inference, width, height, gp=None, batch_size=BATCH_SIZE, conf_thres=CONF_THRES, iou_thres=NMS_IOU, verbose=True):
     """
     Args:
         frames_dir: a directory containing frames to be evaluated
@@ -147,7 +147,7 @@ def do_suppression(dataloader, inference, width, height, gp=None, batch_size=BAT
     # keep predictions to feed them ordered into the Tracker
     # TODO: how to deal with large files?
     all_preds = {}
-    with tqdm(total=len(dataloader)*batch_size, desc="Running suppression", ncols=0) as pbar:
+    with tqdm(total=len(dataloader)*batch_size, desc="Running suppression", ncols=0, disable=not verbose) as pbar:
         for batch_i, (img, _, shapes) in enumerate(dataloader):
 
             if gp: gp(batch_i / len(dataloader), pbar.__str__())
@@ -180,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, max_age=MAX_AGE, iou_thres=IOU_THRES, min_hits=MIN_HITS, min_length=MIN_LENGTH): 
+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, verbose=True): 
 
     if (gp): gp(0, "Tracking...")
 
@@ -194,7 +194,7 @@ def do_tracking(all_preds, image_meter_width, image_meter_height, gp=None, max_a
     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:
+    with tqdm(total=len(all_preds), desc="Running tracking", ncols=0, disable=not verbose) as pbar:
         for i, key in enumerate(sorted(all_preds.keys())):
             if gp: gp(i / len(all_preds), pbar.__str__())
             boxes = all_preds[key]

scripts/infer_frames.py

@@ -8,9 +8,10 @@ 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 tqdm import tqdm
 
 
-def main(args, config={}):
+def main(args, config={}, verbose=True):
     """
     Main processing task to be run in gradio
         - Writes aris frames to dirname(filepath)/frames/{i}.jpg
@@ -54,15 +55,18 @@ def main(args, config={}):
 
         seq_list = os.listdir(in_loc_dir)
         idx = 1
-        for seq in seq_list:
-            print(" ")
-            print("(" + str(idx) + "/" + str(len(seq_list)) + ") " + seq)
-            print(" ")
-            idx += 1
-            in_seq_dir = os.path.join(in_loc_dir, seq)
-            infer_seq(in_seq_dir, out_dir, config, seq, model, device, metadata_path)
-
-def infer_seq(in_dir, out_dir, config, seq_name, model, device, metadata_path):
+        with tqdm(total=len(seq_list), desc="...", ncols=0) as pbar:
+            for seq in seq_list:
+                pbar.update(1)
+                pbar.set_description("Processing " + seq)
+                print(" ")
+                print("(" + str(idx) + "/" + str(len(seq_list)) + ") " + seq)
+                print(" ")
+                idx += 1
+                in_seq_dir = os.path.join(in_loc_dir, seq)
+                infer_seq(in_seq_dir, out_dir, config, seq, model, device, metadata_path, verbose)
+
+def infer_seq(in_dir, out_dir, config, seq_name, model, device, metadata_path, verbose):
     
     #progress_log = lambda p, m: 0
 
@@ -83,16 +87,16 @@ def infer_seq(in_dir, out_dir, config, seq_name, model, device, metadata_path):
     dataloader = create_dataloader_frames_only(in_dir)
 
     try:
-        inference, width, height = do_detection(dataloader, model, device)
+        inference, width, height = do_detection(dataloader, model, device, verbose=verbose)
     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, conf_thres=config['conf_threshold'], iou_thres=config['nms_iou'])
+    all_preds, real_width, real_height = do_suppression(dataloader, inference, width, height, conf_thres=config['conf_threshold'], iou_thres=config['nms_iou'], verbose=verbose)
 
-    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'])
+    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'], verbose=verbose)
 
     mot_rows = []
     for frame in results['frames']: