Merge pull request #20 from naseemap47/tracker

Tracker

app.py

@@ -9,7 +9,6 @@
 import json
 import pandas as pd
 from model_utils import get_yolo, color_picker_fn, get_system_stat
-# from ultralytics import YOLO
 
 
 p_time = 0
@@ -52,7 +51,7 @@
                 model = custom(path_or_model=path_model_file, gpu=True)
 
         # YOLOv8 Model
-        if model_type == 'YOLOv8':
+        elif model_type == 'YOLOv8':
             from ultralytics import YOLO
             model = YOLO(path_model_file)
 
@@ -62,6 +61,8 @@
         # Inference Mode
         options = st.sidebar.radio(
             'Options:', ('Webcam', 'Image', 'Video', 'RTSP'), index=1)
+
+        tracker = st.sidebar.selectbox("Choose Tracker", ("Tracker", "bytetrack", "botsort"))
 
         # Confidence
         confidence = st.sidebar.slider(
@@ -91,7 +92,11 @@
                 FRAME_WINDOW.image(img, channels='BGR')
 
                 if pred:
-                    img, current_no_class = get_yolo(img, model_type, model, confidence, color_pick_list, class_labels, draw_thick)
+                    img, current_no_class = get_yolo(
+                        img, model_type, model, confidence, 
+                        color_pick_list, class_labels, draw_thick,
+                        Tracker=tracker
+                    )
                     FRAME_WINDOW.image(img, channels='BGR')
 
                     # Current number of classes
@@ -107,7 +112,7 @@
                         st.dataframe(df_fq, use_container_width=True)
 
         # Video
-        if options == 'Video':
+        elif options == 'Video':
             upload_video_file = st.sidebar.file_uploader(
                 'Upload Video', type=['mp4', 'avi', 'mkv'])
             if upload_video_file is not None:
@@ -120,7 +125,7 @@
 
 
         # Web-cam
-        if options == 'Webcam':
+        elif options == 'Webcam':
             cam_options = st.sidebar.selectbox('Webcam Channel',
                                             ('Select Channel', '0', '1', '2', '3'))
 
@@ -130,7 +135,7 @@
 
 
         # RTSP
-        if options == 'RTSP':
+        elif options == 'RTSP':
             rtsp_url = st.sidebar.text_input(
                 'RTSP URL:',
                 'eg: rtsp://admin:[email protected]/cam/realmonitor?channel=0&subtype=0'
@@ -152,7 +157,11 @@
             )
             break
 
-        img, current_no_class = get_yolo(img, model_type, model, confidence, color_pick_list, class_labels, draw_thick)
+        img, current_no_class = get_yolo(
+            img, model_type, model, confidence, 
+            color_pick_list, class_labels, draw_thick,
+            Tracker=tracker
+        )
         FRAME_WINDOW.image(img, channels='BGR')
 
         # FPS

model_utils.py

@@ -21,10 +21,10 @@ def color_picker_fn(classname, key):
     return color
 
 
-def get_yolo(img, model_type, model, confidence, color_pick_list, class_list, draw_thick):
+def get_yolo(img, model_type, model, confidence, color_pick_list, class_list, draw_thick, Tracker):
     current_no_class = []
-    results = model(img)
     if model_type == 'YOLOv7':
+        results = model(img)
         box = results.pandas().xyxy[0]
 
         for i in box.index:
@@ -35,18 +35,37 @@ def get_yolo(img, model_type, model, confidence, color_pick_list, class_list, dr
                                 color=color_pick_list[id], line_thickness=draw_thick)
             current_no_class.append([class_name])
 
-    if model_type == 'YOLOv8':
-        for result in results:
-            bboxs = result.boxes.xyxy
-            conf = result.boxes.conf
-            cls = result.boxes.cls
-            for bbox, cnf, cs in zip(bboxs, conf, cls):
-                xmin = int(bbox[0])
-                ymin = int(bbox[1])
-                xmax = int(bbox[2])
-                ymax = int(bbox[3])
-                if cnf > confidence:
-                    plot_one_box([xmin, ymin, xmax, ymax], img, label=class_list[int(cs)],
+    elif model_type == 'YOLOv8':
+        if Tracker == "Tracker":
+            results = model.predict(img, conf=confidence)
+            for result in results:
+                bboxs = result.boxes.xyxy
+                conf = result.boxes.conf
+                cls = result.boxes.cls
+                for bbox, cnf, cs in zip(bboxs, conf, cls):
+                    xmin = int(bbox[0])
+                    ymin = int(bbox[1])
+                    xmax = int(bbox[2])
+                    ymax = int(bbox[3])
+                    # if cnf > confidence:
+                    plot_one_box([xmin, ymin, xmax, ymax], img, label=f"{class_list[int(cs)]} {cnf:.3}",
+                                    color=color_pick_list[int(cs)], line_thickness=draw_thick)
+                    current_no_class.append([class_list[int(cs)]])
+
+        elif Tracker != "Tracker":
+            results = model.track(img, tracker=f"tracker/{Tracker}.yaml")
+            for result in results:
+                bboxs = result.boxes.xyxy
+                conf = result.boxes.conf
+                cls = result.boxes.cls
+                track_id = result.boxes.id.int().cpu().tolist()
+                for bbox, cnf, cs, id in zip(bboxs, conf, cls, track_id):
+                    xmin = int(bbox[0])
+                    ymin = int(bbox[1])
+                    xmax = int(bbox[2])
+                    ymax = int(bbox[3])
+                    # if cnf > confidence:
+                    plot_one_box([xmin, ymin, xmax, ymax], img, label=f"ID: {id} {class_list[int(cs)]} {cnf:.3}",
                                     color=color_pick_list[int(cs)], line_thickness=draw_thick)
                     current_no_class.append([class_list[int(cs)]])
     return img, current_no_class

tracker/botsort.yaml

@@ -0,0 +1,18 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Default YOLO tracker settings for BoT-SORT tracker https://github.com/NirAharon/BoT-SORT
+
+tracker_type: botsort # tracker type, ['botsort', 'bytetrack']
+track_high_thresh: 0.5 # threshold for the first association
+track_low_thresh: 0.1 # threshold for the second association
+new_track_thresh: 0.6 # threshold for init new track if the detection does not match any tracks
+track_buffer: 30 # buffer to calculate the time when to remove tracks
+match_thresh: 0.8 # threshold for matching tracks
+fuse_score: True # Whether to fuse confidence scores with the iou distances before matching
+# min_box_area: 10  # threshold for min box areas(for tracker evaluation, not used for now)
+
+# BoT-SORT settings
+gmc_method: sparseOptFlow # method of global motion compensation
+# ReID model related thresh (not supported yet)
+proximity_thresh: 0.5
+appearance_thresh: 0.25
+with_reid: False

tracker/bytetrack.yaml

@@ -0,0 +1,11 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Default YOLO tracker settings for ByteTrack tracker https://github.com/ifzhang/ByteTrack
+
+tracker_type: bytetrack # tracker type, ['botsort', 'bytetrack']
+track_high_thresh: 0.5 # threshold for the first association
+track_low_thresh: 0.1 # threshold for the second association
+new_track_thresh: 0.6 # threshold for init new track if the detection does not match any tracks
+track_buffer: 30 # buffer to calculate the time when to remove tracks
+match_thresh: 0.8 # threshold for matching tracks
+fuse_score: True # Whether to fuse confidence scores with the iou distances before matching
+# min_box_area: 10  # threshold for min box areas(for tracker evaluation, not used for now)