Update app.py

app.py

@@ -1,177 +1,114 @@
-import os
-
-from gradio_webrtc import WebRTC
-import requests
-from PIL import Image
-
-import matplotlib.pyplot as plt
-
-from random import choice
-import io
-
 import gradio as gr
-
 import cv2
 import numpy as np
-
-from io import BytesIO
-import random
-import tempfile
-from pathlib import Path
-
-import torch
-from transformers import pipeline
-
 from PIL import Image
+import matplotlib.pyplot as plt
+from transformers import pipeline
+import torch
+from random import choice
+from io import BytesIO
+import os
+from datetime import datetime
 
-import matplotlib.patches as patches
-
-
-detector50 = pipeline(model="facebook/detr-resnet-50")
-
-detector101 = pipeline(model="facebook/detr-resnet-101")
-
+# 初始化对象检测器并移动到GPU（如果可用）
+detector = pipeline(model="facebook/detr-resnet-101", use_fast=True)
 if torch.cuda.is_available():
-    # print("##############------------use cuda!------------#################")
-    detector50.model.to('cuda')
-    detector101.model.to('cuda')
-
-model = "detr-resnet-101"
+    detector.model.to('cuda')
 
 COLORS = ["#ff7f7f", "#ff7fbf", "#ff7fff", "#bf7fff",
-            "#7f7fff", "#7fbfff", "#7fffff", "#7fffbf",
-            "#7fff7f", "#bfff7f", "#ffff7f", "#ffbf7f"]
+          "#7f7fff", "#7fbfff", "#7fffff", "#7fffbf",
+          "#7fff7f", "#bfff7f", "#ffff7f", "#ffbf7f"]
 
 fdic = {
-#    "family" : "Impact",
-    "style" : "italic",
-    "size" : 15,
-    "color" : "yellow",
-    "weight" : "bold"
+    "style": "italic",
+    "size": 15,
+    "color": "yellow",
+    "weight": "bold"
 }
 
-#######################################
-
-
-def query_data(model, in_pil_img: Image.Image):
-    results = None
-    if model == "detr-resnet-101":
-        results = detector101(in_pil_img)
-    else:
-        results = detector50(in_pil_img)
-    # print(f"检测结果：{results}")
+def query_data(in_pil_img: Image.Image):
+    results = detector(in_pil_img)
+    print(f"检测结果：{results}")
     return results
 
-
-
-def get_figure(in_pil_img):
+def get_annotated_image(in_pil_img):
     plt.figure(figsize=(16, 10))
     plt.imshow(in_pil_img)
-    
     ax = plt.gca()
-    # print(f"图像尺寸：{in_pil_img.size}")
-    in_results = query_data(model, in_pil_img)
+    in_results = query_data(in_pil_img)
 
     for prediction in in_results:
-        selected_color = choice(COLORS)
-
-        x, y = prediction['box']['xmin'], prediction['box']['ymin'],
-        w, h = prediction['box']['xmax'] - prediction['box']['xmin'], prediction['box']['ymax'] - prediction['box']['ymin']
+        color = choice(COLORS)
+        box = prediction['box']
+        label = prediction['label']
+        score = round(prediction['score'] * 100, 1)
 
-        ax.add_patch(plt.Rectangle((x, y), w, h, fill=False, color=selected_color, linewidth=3))
-        ax.text(x, y, f"{prediction['label']}: {round(prediction['score']*100, 1)}%", fontdict=fdic)
-        # print(f"x: {x}, y: {y}, w: {w}, h: {h}, label: {prediction['label']}, score: {prediction['score']}")
+        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)
 
     plt.axis("off")
-
-    return plt.gcf()
-
-
-def process_single_frame(frame):
-    # print(f"开始处理单帧")
-    # 将 BGR 转换为 RGB
-    rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-    
-    # 创建 PIL 图像对象
-    pil_image = Image.fromarray(rgb_frame)
-
-    # 获取带有标注信息的图像
-    figure = get_figure(pil_image)
-
     buf = BytesIO()
-    figure.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
+    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)
 
+def process_video(input_video_path):
+    cap = cv2.VideoCapture(input_video_path)
+    if not cap.isOpened():
+        raise ValueError("无法打开输入视频文件")
 
-def infer_video(input_video_path):
-    # print(f"开始处理视频 {input_video_path}")
-    with tempfile.TemporaryDirectory() as tmp_dir:
-        # output_video_path = Path(tmp_dir) / "output.mp4"
-        cap = cv2.VideoCapture(input_video_path)
-
-        if not cap.isOpened():
-            raise ValueError("无法打开输入视频文件")
-
-        # width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-        # height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-        # fps = cap.get(cv2.CAP_PROP_FPS)
-        # fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))  # 使用原始视频的编码器
-        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # 获取总帧数
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 使用 'mp4v' 编码器
+    output_dir = './output_videos'  # 指定输出目录
+    os.makedirs(output_dir, exist_ok=True)  # 确保输出目录存在
 
-        # out = cv2.VideoWriter(str(output_video_path), fourcc, fps, (width, height))
+    # 生成唯一文件名
+    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)
 
-        frame_count = 0
-        try:
-            while frame_count < total_frames:
-                ret, frame = cap.read()
-                if not ret:
-                    print(f"提前结束：在第 {frame_count} 帧时无法读取帧")
-                    break
-                
-                frame_count += 1
+    out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
 
-                # 处理单帧并转换为 OpenCV 格式（BGR）
-                processed_frame = process_single_frame(frame)
-                bgr_frame = cv2.cvtColor(processed_frame, cv2.COLOR_RGB2BGR)
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        
+        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        pil_image = Image.fromarray(rgb_frame)
+        annotated_frame = get_annotated_image(pil_image)
+        bgr_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR)
 
-                yield bgr_frame
+        # 确保帧的尺寸与视频输出一致
+        if bgr_frame.shape[:2] != (height, width):
+            bgr_frame = cv2.resize(bgr_frame, (width, height))
 
-                # 可选：显示进度
-                if frame_count % 30 == 0:
-                    print(f"已处理 {frame_count}/{total_frames} 帧")
+        print(f"Writing frame of shape {bgr_frame.shape} and type {bgr_frame.dtype}")  # 调试信息
+        out.write(bgr_frame)
 
-                # if frame_count == 48:
-                #     print("测试结束")
-                #     return None
+    cap.release()
+    out.release()
 
-        finally:
-            cap.release()
-        
-        return None
+    # 返回输出视频路径给 Gradio
+    return output_video_path
 
-    
-# 更新 Gradio 接口以支持视频输入和输出
-with gr.Blocks(title="基于AI的安全风险识别及防控应用",
-               css=".gradio-container {background:lightyellow;}"
-              ) as demo:
+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>")
     
     with gr.Row():
         input_video = gr.Video(label="输入视频")
-        output_video = WebRTC(label="WebRTC Stream",
-                                   rtc_configuration=None,
-                                   mode="receive",
-                                   modality="video")
-        detect = gr.Button("Detect", variant="primary")
-    output_video.stream(
-        fn=infer_video,
-        inputs=[input_video],
-        outputs=[output_video],
-        trigger=detect.click
-    )
-
-demo.launch()
+        detect_button = gr.Button("开始检测", variant="primary")
+        output_video = gr.Video(label="输出视频")
+
+    # 将process_video函数绑定到按钮点击事件，并将处理后的视频路径传递给output_video
+    detect_button.click(process_video, inputs=input_video, outputs=output_video)
 
+demo.launch()