增加对置信度和类别的选择

app.py

@@ -17,20 +17,16 @@ COLORS = ["#ff7f7f", "#ff7fbf", "#ff7fff", "#bf7fff",
           "#7f7fff", "#7fbfff", "#7fffff", "#7fffbf",
           "#7fff7f", "#bfff7f", "#ffff7f", "#ffbf7f"]
 
-fdic = {
-    "style": "italic",
-    "size": 24,
-    "color": "yellow",
-    "weight": "bold"
-}
 
-threshold = 80  # 置信度阈值
+
+threshold = 90  # 置信度阈值
+label_list = ["person", "car", "truck"]
 
 label_color_dict = {}
 
 def query_data(in_pil_img: Image.Image):
     results = detector(in_pil_img)
-    # print(f"检测结果：{results}")
+    print(f"检测结果：{results}")
     return results
 
 
@@ -61,6 +57,8 @@ def get_annotated_image(in_pil_img):
         score = round(prediction['score'] * 100, 1)
         if score < threshold:
             continue  # 过滤掉低置信度的预测结果
+        if label not in label_list:
+            continue  # 过滤掉不在允许显示的label列表中的预测结果
 
         if label not in label_color_dict:  # 为每个类别随机分配颜色, 后续维持一致
             color = choice(COLORS)
@@ -105,7 +103,7 @@ def process_video(input_video_path):
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     output_video_filename = f"output_{timestamp}.mp4"
     output_video_path = os.path.join(output_dir, output_video_filename)
-    # print(f"输出视频信息：{output_video_path}, {width}x{height}, {fps}fps")
+    print(f"输出视频信息：{output_video_path}, {width}x{height}, {fps}fps")
     out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
 
     while True:
@@ -115,15 +113,15 @@ def process_video(input_video_path):
         
         rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
         pil_image = Image.fromarray(rgb_frame)
-        # print(f"Input frame of shape {rgb_frame.shape} and type {rgb_frame.dtype}")  # 调试信息
+        print(f"Input frame of shape {rgb_frame.shape} and type {rgb_frame.dtype}")  # 调试信息
         annotated_frame = get_annotated_image(pil_image)
         bgr_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR)
-        # print(f"Annotated frame of shape {bgr_frame.shape} and type {bgr_frame.dtype}")  # 调试信息
+        print(f"Annotated frame of shape {bgr_frame.shape} and type {bgr_frame.dtype}")  # 调试信息
         # 确保帧的尺寸与视频输出一致
         if bgr_frame.shape[:2] != (height, width):
             bgr_frame = cv2.resize(bgr_frame, (width, height))
 
-        # print(f"Writing frame of shape {bgr_frame.shape} and type {bgr_frame.dtype}")  # 调试信息
+        print(f"Writing frame of shape {bgr_frame.shape} and type {bgr_frame.dtype}")  # 调试信息
         out.write(bgr_frame)
 
     cap.release()
@@ -132,9 +130,31 @@ def process_video(input_video_path):
     # 返回输出视频路径给 Gradio
     return output_video_path
 
+def change_threshold(value):
+    global threshold
+    threshold = value
+    return f"当前置信度阈值为{threshold}%"
+
+def update_labels(selected_labels):
+    # 更新 label_list 以匹配用户的选择
+    global label_list
+    label_list = selected_labels
+    return selected_labels
+
 with gr.Blocks(css=".gradio-container {background:lightyellow;}", title="基于AI的安全风险识别及防控应用") as demo:
     gr.HTML("<div style='font-family:'Times New Roman', 'Serif'; font-size:16pt; font-weight:bold; text-align:center; color:royalblue;'>基于AI的安全风险识别及防控应用</div>")
     
+    # 设置置信度阈值
+    threshold_slider = gr.Slider(minimum=0, maximum=100, value=threshold, step=1, label="置信度阈值")
+    textbox = gr.Textbox(value=f"当前置信度阈值为{threshold}%", label="置信度显示")
+    # 绑定滑块变化事件到change_threshold函数，同时设置输出为textbox
+    threshold_slider.change(fn=change_threshold, inputs=[threshold_slider], outputs=[textbox])
+
+    # 设置允许显示的label列表
+    label_checkboxes = gr.CheckboxGroup(choices=label_list, value=label_list, label="检测目标")
+    # 允许修改label_list
+    label_checkboxes.change(fn=update_labels, inputs=[label_checkboxes], outputs=[label_checkboxes])
+
     with gr.Row():
         input_video = gr.Video(label="输入视频")
         detect_button = gr.Button("开始检测", variant="primary")