switch pretrained model

app.py

@@ -1,70 +1,70 @@
-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()
+import gradio as gr
+import supervision as sv
+from func import detect_and_track
+from transformers import AutoImageProcessor, AutoModelForObjectDetection
+
+processor = AutoImageProcessor.from_pretrained("PekingU/rtdetr_r50vd_coco_o365")
+model = AutoModelForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd_coco_o365")
+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()