diff --git a/app.py b/app.py index aa4c3f2..643dc66 100644 --- a/app.py +++ b/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:name6666@198.162.1.58/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 diff --git a/model_utils.py b/model_utils.py index 53396de..a9324c1 100644 --- a/model_utils.py +++ b/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 diff --git a/tracker/botsort.yaml b/tracker/botsort.yaml new file mode 100644 index 0000000..3683044 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tracker/bytetrack.yaml b/tracker/bytetrack.yaml new file mode 100644 index 0000000..3465b2e --- /dev/null +++ b/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) \ No newline at end of file