Update app.py

app.py

@@ -1,12 +1,10 @@
 import gradio as gr
 import cv2
 import numpy as np
-from PIL import Image
-import matplotlib.pyplot as plt
+from PIL import Image,ImageDraw
 from transformers import pipeline
 import torch
 from random import choice
-from io import BytesIO
 import os
 from datetime import datetime
 
@@ -21,41 +19,43 @@ COLORS = ["#ff7f7f", "#ff7fbf", "#ff7fff", "#bf7fff",
 
 fdic = {
     "style": "italic",
-    "size": 15,
+    "size": 16,
     "color": "yellow",
     "weight": "bold"
 }
 
+label_color_dict = {}
+
 def query_data(in_pil_img: Image.Image):
     results = detector(in_pil_img)
     # print(f"检测结果：{results}")
     return results
 
 def get_annotated_image(in_pil_img):
-    plt.figure(figsize=(16, 10))
-    plt.imshow(in_pil_img)
-    ax = plt.gca()
+    draw = ImageDraw.Draw(in_pil_img)
     in_results = query_data(in_pil_img)
 
     for prediction in in_results:
-        color = choice(COLORS)
         box = prediction['box']
         label = prediction['label']
         score = round(prediction['score'] * 100, 1)
+        if score < 50:
+            continue # 过滤掉低置信度的预测结果
+
+        if label not in label_color_dict: # 为每个类别随机分配颜色, 后续维持一致
+            color = choice(COLORS)
+            label_color_dict[label] = color
+        else:
+            color = label_color_dict[label]
 
-        ax.add_patch(plt.Rectangle((box['xmin'], box['ymin']), 
-                                   box['xmax'] - box['xmin'],
-                                   box['ymax'] - box['ymin'],
-                                   fill=False, color=color, linewidth=3))
-        ax.text(box['xmin'], box['ymin'], f"{label}: {score}%", fontdict=fdic)
+        # 绘制矩形
+        draw.rectangle([box['xmin'], box['ymin'], box['xmax'], box['ymax']], outline=color, width=3)
+        
+        # 添加文本
+        draw.text((box['xmin'], box['ymin']), f"{label}: {score}%", fill=color, fontdict=fdic)
 
-    plt.axis("off")
-    buf = BytesIO()
-    plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
-    plt.close()  # 关闭图形以释放内存
-    buf.seek(0)
-    annotated_image = Image.open(buf).convert('RGB')
-    return np.array(annotated_image)
+    # 返回的是原始图像对象，它已经被修改了
+    return np.array(in_pil_img.convert('RGB'))
 
 def process_video(input_video_path):
     cap = cv2.VideoCapture(input_video_path)
@@ -74,7 +74,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")
     out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
 
     while True:
@@ -84,9 +84,10 @@ 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}")  # 调试信息
         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}")  # 调试信息
         # 确保帧的尺寸与视频输出一致
         if bgr_frame.shape[:2] != (height, width):
             bgr_frame = cv2.resize(bgr_frame, (width, height))