File size: 2,123 Bytes
ff18d07 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import gradio as gr
import supervision as sv
from func import detect_and_track
from transformers import DetrForObjectDetection, DetrImageProcessor
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
tracker = sv.ByteTrack()
mask_annotator = sv.MaskAnnotator()
bbox_annotator = sv.BoundingBoxAnnotator()
label_annotator = sv.LabelAnnotator()
def process_video(video_path, confidence_threshold):
return detect_and_track(
video_path,
model,
processor,
tracker,
confidence_threshold,
mask_annotator,
bbox_annotator,
label_annotator,
)
with gr.Blocks() as demo:
with gr.Row():
with gr.Column():
in_video = gr.Video(
label="待检测视频",
show_download_button=True,
show_share_button=True,
)
slide_cofidence = gr.Slider(
minimum=0.0, maximum=1.0, value=0.8, label="置信度阈值"
)
examples = gr.Examples(
examples=[
"./demo_video/blurry.mp4",
"./demo_video/high-way.mp4",
"./demo_video/aerial.mp4",
],
inputs=in_video,
label="案例视频",
)
with gr.Column():
out_video = gr.Video(
label="检测结果视频",
interactive=False,
show_download_button=True,
show_share_button=True,
)
combine_video = gr.Video(
interactive=False,
label="前后对比",
show_download_button=True,
show_share_button=True,
)
start_detect = gr.Button(value="开始检测")
start_detect.click(
fn=process_video,
inputs=[in_video, slide_cofidence],
outputs=[out_video, combine_video],
)
demo.launch()
|