projects/human_detection/app.py

import os

from models.detectors.yolov7 import YOLOv7ONNX,YOLOv7TRT
from models.detectors.mmyolov8 import MMYOLOv8ONNX,MMYOLOv8TRT
from projects.human_detection.engine.pipeline import run_e2e_pipeline
from mmcv import VideoReader
import gradio as gr

title = "Human Detection - CYBERCORE AI DEMO"
description = "Human Monitoring Demo: It will run detection, tracking on input video and show output video.\nYou may click on of the examples or upload your own image."

root_folder = "/data/human_detection"
example_video_paths = [
    os.path.join(root_folder,"sample_inputs/1.mp4"),
    os.path.join(root_folder,"sample_inputs/2.mp4"),
    os.path.join(root_folder,"sample_inputs/3.mp4"),
]
output_dir = os.path.join(root_folder,"outputs")
os.makedirs(output_dir, exist_ok=True)
assert any([os.path.exists(example_video_path) for example_video_path in example_video_paths]), f"Example video does not exist. Please download example videos from NAS:[https://gofile.me/6ZWyr/q0saLr8hb] and put them in the following path 'human_detection/inputs'"


#---------------------Model Path----------------------------------
# model_path_yolov7_onnx = os.path.join(root_folder, "weights/yolov7_pedestrian_480x640.onnx") # model_path yolov7
# model_path_yolov7_trt  = os.path.join(root_folder, "weights/yolov7_pedestrian_480x640.trt") # model_path yolov7
# assert os.path.exists(model_path_yolov7_onnx) or os.path.exists(model_path_yolov7_trt), f"Model does not exist. Please download model from NAS:[path], compile TRT, and put it in the following path 'project_folder'/weights/yolov7-honda-pedestrian_deploy_v2_0.1_768x1280.onnx"
model_path_yolov8_onnx = os.path.join(root_folder, "weights/mmyolov8_s_human_dynamic_shape.onnx") # model_path yolov8
model_path_yolov8_trt  = os.path.join(root_folder, "weights/mmyolov8_s_human_DBsx640x800.trt") # model_path yolov8
assert os.path.exists(model_path_yolov8_onnx) or os.path.exists(model_path_yolov8_trt), f"Model does not exist. Please download model from NAS:[path], compile TRT, and put it in the following path 'project_folder'/weights/yolov8-human.onnx"

# ---------------------Configs---------------------------------- 
use_trt_yolov8 = os.path.exists(model_path_yolov8_trt)
yolov8_cfg = dict(
    img_shape=(3, 640, 800),
    batch_size=32,
    preprocess_cfg=dict(
        border_color=(114, 114, 114),
        auto=False,
        scaleFill=False,
        scaleup=False,
        stride=32),
    nms_agnostic_cfg=dict(
        type='nms',
        iou_threshold=0.6,
        class_agnostic=True),
    score_thr=0.1,
    model_path = model_path_yolov8_trt if use_trt_yolov8 else model_path_yolov8_onnx,
    device='0'
)

# use_trt_yolov7 = os.path.exists(model_path_yolov7_trt)
# yolov7_cfg = dict(
#     img_shape=(3, 480, 640),
#     batch_size=32,
#     preprocess_cfg=dict(
#         border_color=(114, 114, 114),
#         auto=False,
#         scaleFill=False,
#         scaleup=True,
#         stride=32),
#     nms_agnostic_cfg=dict(
#         type='nms',
#         iou_threshold=0.99,
#         class_agnostic=True),
#     model_path = model_path_yolov7_trt if use_trt_yolov7 else model_path_yolov7_onnx,
#     device='0'
# )

tracker_cfg = dict(
    obj_score_thrs=dict(high=0.6, low=0.1),
    init_track_thr=0.7,
    weight_iou_with_det_scores=True,
    match_iou_thrs=dict(high=0.1, low=0.5, tentative=0.3),
    num_frames_retain=30,
    motion=dict(type='KalmanFilter')
)

visualizer_cfg = dict(fps=-1, min_width=1280)


# yolov7_detector = YOLOv7TRT(**yolov7_cfg) if use_trt_yolov7 else YOLOv7ONNX(**yolov7_cfg)
yolov8_detector = MMYOLOv8TRT(**yolov8_cfg) if use_trt_yolov8 else MMYOLOv8ONNX(**yolov8_cfg)
    
def inference(video, model_name, conf_thres, show_conf, progress=gr.Progress()):
    output_path = os.path.join(output_dir, os.path.basename(video))
    if os.path.exists(output_path):
        os.remove(output_path)
    # detector = yolov7_detector if model_name == "pedestrian" else yolov8_detector
    detector = yolov8_detector
    run_e2e_pipeline(video, detector, tracker_cfg, visualizer_cfg, output_path, conf_thres, show_conf, progress)
    return output_path

def clear(*all_components):
    outputs = [None]*len(all_components)
    return outputs

# ---------------------Gradio UI----------------------------------
with gr.Blocks(title=title) as demo:
    gr.Markdown("<h1 style='text-align: center; margin-bottom: 1rem'>" + title + "</h1>")
    gr.Markdown(description)
    input_components = []
    output_components = []

    with gr.Row():
        input_video = gr.Video(type="file", label="input_video")
        output_video = gr.Video(label="output_video")
        input_components.append(input_video)
        output_components.append(output_video)

    with gr.Row().style(equal_height=True, mobile_collapse=True):
        with gr.Column(scale=2, variant="panel") as input_column:
            model_dropdown = gr.Dropdown(label="Detector Model", 
                                         choices=["pedestrian", "general-human"], 
                                         default="general-human", 
                                         info="Choose the application model for Human detection.")
            prob_threshold_slider = gr.components.Slider(minimum=0, maximum=1.0, step=0.01, value=0.3, label="Confidence Threshold")
            show_confidence = gr.Checkbox(label="Show Confidence")
            input_components.extend([model_dropdown, prob_threshold_slider, show_confidence])
        
        with gr.Column(scale=2):
            examples_handler = gr.Examples(
                examples=[[item] for item in example_video_paths],
                fn=inference,
                inputs=input_components,
                outputs=output_components,
                examples_per_page=3
            )
            
            with gr.Row():
                submit_btn = gr.Button("Submit", variant="primary")
                clear_btn = gr.Button("Clear")
    
    submit_btn.click(
        inference,
        input_components,
        output_components,
        api_name="predict",
        scroll_to_output=True,
    )

    clear_btn.click(
        clear,
        input_components + output_components,
        input_components + output_components,
    )

demo.queue(concurrency_count=3).launch(share=True, server_name='0.0.0.0', server_port=7860)