Tracking script

scripts/detect_frames.py

@@ -12,7 +12,7 @@ from tqdm import tqdm
 import numpy as np
 
 
-def main(args, config={}, verbose=True):
+def main(args, verbose=True):
     """
     Main processing task to be run in gradio
         - Writes aris frames to dirname(filepath)/frames/{i}.jpg
@@ -28,16 +28,6 @@ def main(args, config={}, verbose=True):
     print("In task...")
     print("Cuda available in task?", torch.cuda.is_available())
 
-    # setup config
-    if "conf_threshold" not in config: config['conf_threshold'] = 0.001
-    if "nms_iou" not in config: config['nms_iou'] = 0.6
-    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)
-
     model, device = setup_model(args.weights)
     
     locations = [
@@ -50,11 +40,11 @@ def main(args, config={}, verbose=True):
         print(in_loc_dir)
         print(out_loc_dir)
 
-        detect_location(in_loc_dir, out_loc_dir, config, model, device, verbose)
+        detect_location(in_loc_dir, out_loc_dir, model, device, verbose)
 
 
                 
-def detect_location(in_loc_dir, out_loc_dir, config, model, device, verbose):
+def detect_location(in_loc_dir, out_loc_dir, model, device, verbose):
 
     seq_list = os.listdir(in_loc_dir)
 
@@ -69,63 +59,25 @@ def detect_location(in_loc_dir, out_loc_dir, config, model, device, verbose):
             out_seq_dir = os.path.join(out_loc_dir, seq)
             os.makedirs(out_seq_dir, exist_ok=True)
 
-            detect_seq(in_seq_dir, out_seq_dir, config, model, device, verbose)
-
-def detect_seq(in_seq_dir, out_seq_dir, config, model, device, verbose):
-
-    ann_list = []
-    frame_list = detect(in_seq_dir, config, model, device, verbose)
-    for frame in frame_list:
-        if frame is not None:
-            for ann in frame:
-                ann_list.append({
-                    'image_id': ann[5],
-                    'category_id': 0,
-                    'bbox': [ann[0], ann[1], ann[2] - ann[0], ann[3] - ann[1]],
-                    'score': ann[4]
-                })
-    result = json.dumps(ann_list)
+            detect(in_seq_dir, out_seq_dir, model, device, verbose)
 
-    with open(os.path.join(out_seq_dir, 'pred.json'), 'w') as f:
-        f.write(result)
-
-def detect(in_dir, config, model, device, verbose):
-    
-    #progress_log = lambda p, m: 0
+def detect(in_seq_dir, out_seq_dir, model, device, verbose):
 
     # create dataloader
-    dataloader = create_dataloader_frames_only(in_dir)
+    dataloader = create_dataloader_frames_only(in_seq_dir)
 
     inference, image_shapes, width, height = do_detection(dataloader, model, device, verbose=verbose)
 
+    json_obj = {
+        'inference': inference,
+        'image_shapes': image_shapes,
+        'width': width,
+        'height': height
+    }
+    json_string = json.dumps(json_obj)
 
-    outputs = do_suppression(inference, conf_thres=config['conf_threshold'], iou_thres=config['nms_iou'], verbose=verbose)
-
-    file_names = dataloader.files
-    frame_list = []
-    for batch_i, batch in enumerate(outputs):
-
-        batch_shapes = image_shapes[batch_i]
-
-        # Format results
-        for si, pred in enumerate(batch):
-            (image_shape, original_shape) = batch_shapes[si]
-
-            # Clip boxes to image bounds and resize to input shape
-            clip_boxes(pred, (height, width))
-            boxes = pred[:, :4].clone()  # xyxy
-            confs = pred[:, 4].clone().tolist()
-            scale_boxes(image_shape, boxes, original_shape[0], original_shape[1])  # to original shape
-
-            frame = [ [*bb, conf] for bb, conf in zip(boxes.tolist(), confs) ]
-            
-            file_name = file_names[batch_i*32 + si]
-            for ann in frame: 
-                ann.append(file_name)
-
-            frame_list.append(frame)
-            
-    return frame_list
+    with open(os.path.join(out_seq_dir, 'pred.json'), 'w') as f:
+        f.write(json_string)
 
 def argument_parser():
     parser = argparse.ArgumentParser()

scripts/full_detect_frames.py

@@ -0,0 +1,139 @@
+import project_path
+from lib.yolov5.utils.general import clip_boxes, scale_boxes
+import argparse
+from datetime import datetime
+import torch
+import os
+from dataloader import create_dataloader_frames_only
+from inference import setup_model, do_detection, do_suppression, do_confidence_boost, format_predictions, do_tracking
+from visualizer import generate_video_batches
+import json
+from tqdm import tqdm
+import numpy as np
+
+
+def main(args, config={}, verbose=True):
+    """
+    Main processing task to be run in gradio
+        - Writes aris frames to dirname(filepath)/frames/{i}.jpg
+        - Writes json output to dirname(filepath)/{filename}_results.json
+        - Writes manual marking to dirname(filepath)/{filename}_marking.txt
+        - Writes video output to dirname(filepath)/{filename}_results.mp4
+        - Zips all results to dirname(filepath)/{filename}_results.zip
+    Args:
+        filepath (str): path to aris file
+        
+    TODO: Separate into subtasks in different queues; have a GPU-only queue.
+    """
+    print("In task...")
+    print("Cuda available in task?", torch.cuda.is_available())
+
+    # setup config
+    if "conf_threshold" not in config: config['conf_threshold'] = 0.001
+    if "nms_iou" not in config: config['nms_iou'] = 0.6
+    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)
+
+    model, device = setup_model(args.weights)
+    
+    locations = [
+        "kenai-val"
+    ]
+    for loc in locations:
+
+        in_loc_dir = os.path.join(args.frames, loc)
+        out_loc_dir = os.path.join(args.output, loc)
+        print(in_loc_dir)
+        print(out_loc_dir)
+
+        detect_location(in_loc_dir, out_loc_dir, config, model, device, verbose)
+
+
+                
+def detect_location(in_loc_dir, out_loc_dir, config, model, device, verbose):
+
+    seq_list = os.listdir(in_loc_dir)
+
+    with tqdm(total=len(seq_list), desc="...", ncols=0) as pbar:
+        for seq in seq_list:
+
+            pbar.update(1)
+            if (seq.startswith(".")): continue
+            pbar.set_description("Processing " + seq)
+
+            in_seq_dir = os.path.join(in_loc_dir, seq)
+            out_seq_dir = os.path.join(out_loc_dir, seq)
+            os.makedirs(out_seq_dir, exist_ok=True)
+
+            detect_seq(in_seq_dir, out_seq_dir, config, model, device, verbose)
+
+def detect_seq(in_seq_dir, out_seq_dir, config, model, device, verbose):
+
+    ann_list = []
+    frame_list = detect(in_seq_dir, config, model, device, verbose)
+    for frame in frame_list:
+        if frame is not None:
+            for ann in frame:
+                ann_list.append({
+                    'image_id': ann[5],
+                    'category_id': 0,
+                    'bbox': [ann[0], ann[1], ann[2] - ann[0], ann[3] - ann[1]],
+                    'score': ann[4]
+                })
+    result = json.dumps(ann_list)
+
+    with open(os.path.join(out_seq_dir, 'pred.json'), 'w') as f:
+        f.write(result)
+
+def detect(in_dir, config, model, device, verbose):
+    
+    #progress_log = lambda p, m: 0
+
+    # create dataloader
+    dataloader = create_dataloader_frames_only(in_dir)
+
+    inference, image_shapes, width, height = do_detection(dataloader, model, device, verbose=verbose)
+
+
+    outputs = do_suppression(inference, conf_thres=config['conf_threshold'], iou_thres=config['nms_iou'], verbose=verbose)
+
+    file_names = dataloader.files
+    frame_list = []
+    for batch_i, batch in enumerate(outputs):
+
+        batch_shapes = image_shapes[batch_i]
+
+        # Format results
+        for si, pred in enumerate(batch):
+            (image_shape, original_shape) = batch_shapes[si]
+
+            # Clip boxes to image bounds and resize to input shape
+            clip_boxes(pred, (height, width))
+            boxes = pred[:, :4].clone()  # xyxy
+            confs = pred[:, 4].clone().tolist()
+            scale_boxes(image_shape, boxes, original_shape[0], original_shape[1])  # to original shape
+
+            frame = [ [*bb, conf] for bb, conf in zip(boxes.tolist(), confs) ]
+            
+            file_name = file_names[batch_i*32 + si]
+            for ann in frame: 
+                ann.append(file_name)
+
+            frame_list.append(frame)
+            
+    return frame_list
+
+def argument_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--frames", required=True, help="Path to frame directory. Required.")
+    parser.add_argument("--output", required=True, help="Path to output directory. 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
+
+if __name__ == "__main__":
+    args = argument_parser().parse_args()
+    main(args)

scripts/track_detection.py

@@ -0,0 +1,141 @@
+import project_path
+from lib.yolov5.utils.general import clip_boxes, scale_boxes
+import argparse
+from datetime import datetime
+import torch
+import os
+from dataloader import create_dataloader_frames_only
+from inference import setup_model, do_detection, do_suppression, do_confidence_boost, format_predictions, do_tracking
+from visualizer import generate_video_batches
+import json
+from tqdm import tqdm
+import numpy as np
+
+
+def main(args, config={}, verbose=True):
+    """
+    Main processing task to be run in gradio
+        - Writes aris frames to dirname(filepath)/frames/{i}.jpg
+        - Writes json output to dirname(filepath)/{filename}_results.json
+        - Writes manual marking to dirname(filepath)/{filename}_marking.txt
+        - Writes video output to dirname(filepath)/{filename}_results.mp4
+        - Zips all results to dirname(filepath)/{filename}_results.zip
+    Args:
+        filepath (str): path to aris file
+        
+    TODO: Separate into subtasks in different queues; have a GPU-only queue.
+    """
+
+    # setup config
+    if "conf_threshold" not in config: config['conf_threshold'] = 0.001
+    if "nms_iou" not in config: config['nms_iou'] = 0.6
+    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
+    if "use_associative" not in config: config['use_associative'] = False
+
+    print(config)
+
+    
+    locations = [
+        "kenai-val"
+    ]
+    for loc in locations:
+
+        in_loc_dir = os.path.join(args.detections, loc)
+        out_loc_dir = os.path.join(args.output, "tracker", "data", loc)
+        os.makedirs(out_loc_dir, exist_ok=True)
+        metadata_path = os.path.join(args.metadata, loc + ".json")
+        print(in_loc_dir)
+        print(out_loc_dir)
+        print(metadata_path)
+
+        track_location(in_loc_dir, out_loc_dir, metadata_path, config, verbose)
+
+
+                
+def track_location(in_loc_dir, out_loc_dir, metadata_path, config, verbose):
+
+    seq_list = os.listdir(in_loc_dir)
+
+    with tqdm(total=len(seq_list), desc="...", ncols=0) as pbar:
+        for seq in seq_list:
+
+            pbar.update(1)
+            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)
+
+def track(in_path, out_path, metadata_path, seq_name, config, verbose):
+
+    # read detection
+    with open(in_path, 'r') as f:
+        detection = json.read(f)
+
+    inference = detection['inference']
+    image_shapes = detection['image_shapes']
+    width = detection['width']
+    height = detection['height']
+
+    # read metadata
+    image_meter_width = -1
+    image_meter_height = -1
+    with open(metadata_path, 'r') as f:
+        json_object = json.loads(f.read())
+        for seq in json_object:
+            if seq['clip_name'] == seq_name:
+                image_meter_width = seq['x_meter_stop'] - seq['x_meter_start']
+                image_meter_height = seq['y_meter_stop'] - seq['y_meter_start']
+
+    outputs = do_suppression(inference, conf_thres=config['conf_threshold'], iou_thres=config['nms_iou'], verbose=verbose)
+
+    if config['use_associative']:
+        do_confidence_boost(inference, outputs, verbose=verbose)
+
+        outputs = do_suppression(inference, conf_thres=config['conf_threshold'], iou_thres=config['nms_iou'], verbose=verbose)
+
+    all_preds, real_width, real_height = format_predictions(image_shapes, outputs, width, height)
+
+    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']:
+        for fish in frame['fish']:
+            bbox = fish['bbox']
+            row = []
+            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(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")
+            row.append("-1")
+            mot_rows.append(",".join(row))
+
+    mot_text = "\n".join(mot_rows)
+
+    with open(out_path, 'w') as f:
+        f.write(mot_text)
+
+def argument_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--detections", required=True, help="Path to frame directory. Required.")
+    parser.add_argument("--output", required=True, help="Path to output directory. Required.")
+    return parser
+
+if __name__ == "__main__":
+    args = argument_parser().parse_args()
+    main(args)

scripts/track_eval.py

@@ -0,0 +1,47 @@
+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 = "../frames/metadata"
+    infer_args.frames = "../frames/images"
+    infer_args.output = "../frames/result"
+    infer_args.weights = "models/v5m_896_300best.pt"
+
+    config = {
+        'conf_threshold': float(args.conf_threshold),
+        'nms_iou': float(args.nms_iou),
+        'min_length': float(args.min_length),
+        'max_age': int(args.max_age),
+        'iou_threshold': float(args.iou_threshold),
+        'min_hits': int(args.min_hits)
+    }
+
+    infer(infer_args, config=config, verbose=False)
+
+    evaluate("../frames/result_testing", "../frames/MOT", "../frames/metadata", "tracker", False)
+
+
+def argument_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--conf_threshold", default=0.3, help="Config object. Required.")
+    parser.add_argument("--nms_iou", default=0.3, help="Config object. Required.")
+    parser.add_argument("--min_length", default=0.3, help="Config object. Required.")
+    parser.add_argument("--max_age", default=20, help="Config object. Required.")
+    parser.add_argument("--iou_threshold", default=0.01, help="Config object. Required.")
+    parser.add_argument("--min_hits", default=11, help="Config object. Required.")
+    return parser
+
+if __name__ == "__main__":
+    args = argument_parser().parse_args()
+    main(args)