Update app.py

app.py

@@ -1,107 +1,79 @@
 import gradio as gr
 import cv2
-import tempfile
+import numpy as np
+import os
 from ultralytics import YOLO
-from pathlib import Path
 
-# 全局变量存储当前模型
-current_model = None
+# 设置上传和结果文件夹
+UPLOAD_FOLDER = 'uploads'
+RESULT_FOLDER = 'results'
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+os.makedirs(RESULT_FOLDER, exist_ok=True)
 
-def load_model(model_path):
-    global current_model
-    try:
-        current_model = YOLO(model_path)
-        return "模型加载成功！"
-    except Exception as e:
-        return f"模型加载失败：{str(e)}"
+# 加载模型
+model = YOLO('yolov8n.pt')
 
-def detect_image(input_image, conf_threshold):
-    if current_model is None:
-        raise gr.Error("请先上传模型文件")
-    
-    results = current_model(input_image, conf=conf_threshold)
-    plotted = results[0].plot()
-    return plotted[:, :, ::-1]  # BGR转RGB
+def process_image(image):
+    # 保存上传的图像
+    filename = 'uploaded_image.jpg'
+    file_path = os.path.join(UPLOAD_FOLDER, filename)
+    cv2.imwrite(file_path, image)
 
-def detect_video(input_video, conf_threshold):
-    if current_model is None:
-        raise gr.Error("请先上传模型文件")
-    
-    cap = cv2.VideoCapture(input_video)
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    
-    # 创建临时输出文件
-    temp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
-    out = cv2.VideoWriter(temp_file.name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
-    
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-        
-        results = current_model(frame, conf=conf_threshold)
-        plotted = results[0].plot()
-        out.write(plotted)
-    
-    cap.release()
-    out.release()
-    return temp_file.name
+    # 处理图像
+    results = model(image)
+    detection_results = []
+    class_counts = {}
 
-def detect_webcam(camera_input, conf_threshold):
-    if current_model is None:
-        raise gr.Error("请先上传模型文件")
-    
-    if camera_input is None:
-        return None
-    
-    results = current_model(camera_input, conf=conf_threshold)
-    plotted = results[0].plot()
-    return plotted[:, :, ::-1]  # BGR转RGB
+    for result in results:
+        boxes = result.boxes
+        for box in boxes:
+            x1, y1, x2, y2 = box.xyxy[0]
+            conf = box.conf[0]
+            cls = box.cls[0]
+            class_name = model.names[int(cls)]
 
-with gr.Blocks() as demo:
-    gr.Markdown("# YOLOv8 自定义模型检测系统")
-    
-    with gr.Row():
-        model_input = gr.File(label="上传模型文件 (.pt)", type="filepath")
-        model_status = gr.Textbox(label="模型状态", interactive=False)
-    
-    model_input.upload(fn=load_model, inputs=model_input, outputs=model_status)
-    
-    with gr.Tabs():
-        with gr.TabItem("图片检测"):
-            with gr.Row():
-                img_input = gr.Image(label="输入图片", type="filepath")
-                img_output = gr.Image(label="检测结果")
-            img_conf = gr.Slider(0, 1, value=0.5, label="置信度阈值")
-            img_button = gr.Button("执行检测")
-            # 正确缩进位置 ↓
-            img_button.click(fn=detect_image, inputs=[img_input, img_conf], outputs=img_output)
-        
-        with gr.TabItem("视频检测"):
-            with gr.Row():
-                video_input = gr.Video(label="输入视频")
-                video_output = gr.Video(label="检测结果")
-            video_conf = gr.Slider(0, 1, value=0.5, label="置信度阈值")
-            video_button = gr.Button("执行检测")
-            # 正确缩进位置 ↓
-            video_button.click(fn=detect_video, inputs=[video_input, video_conf], outputs=video_output)
-        
-        with gr.TabItem("实时摄像头"):
-            webcam_input = gr.Webcam(label="摄像头画面")
-            webcam_output = gr.Image(label="检测结果")
-            webcam_conf = gr.Slider(0, 1, value=0.5, label="置信度阈值")
-            webcam_button = gr.Button("开始检测")
-            # 正确缩进位置 ↓
-            webcam_button.click(fn=detect_webcam, inputs=[webcam_input, webcam_conf], outputs=webcam_output)
+            # 计算每种类别的数量
+            if class_name in class_counts:
+                class_counts[class_name] += 1
+            else:
+                class_counts[class_name] = 1
 
-if __name__ == "__main__":
-    demo.launch()
+            # 绘制边界框和标签
+            cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
+            cv2.putText(image, f'{class_name}:{conf:.2f}', (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9,
+                        (36, 255, 12), 2)
+            detection_results.append(f'Class: {class_name}, Confidence: {conf:.2f}, Box: ({x1}, {y1}), ({x2}, {y2})')
 
-    webcam_button.click(fn=detect_webcam, inputs=[webcam_input, webcam_conf], outputs=webcam_output)
-    
-    # 绑定事件处理
-    img_button.click(fn=detect_image, inputs=[img_input, img_conf], outputs=img_output)
-    video_button.click(fn=detect_video, inputs=[video_input, video_conf], outputs=video_output)
-    webcam_button.click(fn=detect_webcam, inputs=[webcam_input, webcam_conf], outputs=webcam_output)
+    # 在图像上显示检测到的物体数量信息
+    y_offset = 30
+    for class_name, count in class_counts.items():
+        cv2.putText(image, f'{class_name}: {count}', (10, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
+        y_offset += 30
+
+    # 保存处理后的图像
+    result_filename = 'result_image.jpg'
+    result_path = os.path.join(RESULT_FOLDER, result_filename)
+    cv2.imwrite(result_path, image)
+
+    return image, '\n'.join(detection_results)
+
+# 创建Gradio界面
+iface = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(type="numpy", label="上传图像"),
+    outputs=[gr.Image(type="numpy", label="处理后的图像"), gr.Textbox(label="检测结果")],
+    title="YOLOv8 图像检测",
+    description="上传图像并使用YOLOv8模型进行检测"
+)
+
+# 启动Gradio应用
+iface.launch(
+    share=True,
+    server_name="0.0.0.0",
+    server_port=7860,
+    debug=True,
+    #auth=("username", "password"),
+    auth=("test", "12345"),
+    auth_message="Please enter your credentials to access the app.",
+    inbrowser=True
+)