Create app.py

app.py

@@ -0,0 +1,70 @@
+import gradio as gr
+import cv2
+import numpy as np
+from ultralytics import YOLO
+
+# 加载YOLO模型
+model = YOLO('yolo11s-earth.pt')  # 加载你的模型
+
+# 默认类别
+default_classes = [
+    'airplane', 'airport', 'baseballfield', 'basketballcourt', 'bridge',
+    'chimney', 'dam', 'Expressway-Service-area', 'Expressway-toll-station',
+    'golffield', 'groundtrackfield', 'harbor', 'overpass', 'ship',
+    'stadium', 'storagetank', 'tenniscourt', 'trainstation', 'vehicle',
+    'windmill'
+]
+
+def process_frame(frame, classes_input):
+    # 将输入的类别字符串转为列表
+    if classes_input:
+        classes_list = [cls.strip() for cls in classes_input.split(',')]
+        model.set_classes(classes_list)  # 设置模型的类别
+    else:
+        # 如果没有输入，则使用默认类别
+        model.set_classes(default_classes)
+    
+    # 复制帧为可写数组
+    frame = frame.copy()
+    
+    # 转换图像格式
+    rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    
+    # 使用模型进行检测
+    results = model.predict(rgb_frame)
+    
+    # 绘制检测结果
+    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)]
+            
+            # 绘制边界框和标签
+            cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
+            cv2.putText(frame, f'{class_name}:{conf:.2f}', (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2)
+    
+    return frame
+
+def main():
+    # 创建Gradio界面
+    with gr.Blocks() as demo:
+        gr.Markdown("# YOLOv11s-Earth 实时检测")
+        with gr.Row():
+            cam_input = gr.Image(type="numpy", sources=["webcam"], streaming=True, label="摄像头输入")
+            classes_input = gr.Textbox(label="输入类别（逗号分隔）", placeholder="例如：0,1,2")
+        output = gr.Image(label="检测结果", type="numpy")
+        
+        cam_input.stream(
+            process_frame,
+            inputs=[cam_input, classes_input],
+            outputs=output
+        )
+    
+    # 启动Gradio应用
+    demo.launch()
+
+if __name__ == "__main__":
+    main()