Update app.py

app.py

@@ -1,79 +1,43 @@
 import gradio as gr
-import cv2
-import numpy as np
-import os
-from ultralytics import YOLO
+import PIL.Image as Image
+from ultralytics import ASSETS, YOLO
 
-# 设置上传和结果文件夹
-UPLOAD_FOLDER = 'uploads'
-RESULT_FOLDER = 'results'
-os.makedirs(UPLOAD_FOLDER, exist_ok=True)
-os.makedirs(RESULT_FOLDER, exist_ok=True)
+model = YOLO("yolov8n.pt")
 
-# 加载模型
-model = YOLO('yolov8n.pt')
 
-def process_image(image):
-    # 保存上传的图像
-    filename = 'uploaded_image.jpg'
-    file_path = os.path.join(UPLOAD_FOLDER, filename)
-    cv2.imwrite(file_path, image)
+def predict_image(img, conf_threshold, iou_threshold):
+    """Predicts and plots labeled objects in an image using YOLOv8 model with adjustable confidence and IOU thresholds."""
+    results = model.predict(
+        source=img,
+        conf=conf_threshold,
+        iou=iou_threshold,
+        show_labels=True,
+        show_conf=True,
+        imgsz=640,
+    )
 
-    # 处理图像
-    results = model(image)
-    detection_results = []
-    class_counts = {}
+    for r in results:
+        im_array = r.plot()
+        im = Image.fromarray(im_array[..., ::-1])
 
-    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)]
+    return im
 
-            # 计算每种类别的数量
-            if class_name in class_counts:
-                class_counts[class_name] += 1
-            else:
-                class_counts[class_name] = 1
 
-            # 绘制边界框和标签
-            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})')
-
-    # 在图像上显示检测到的物体数量信息
-    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模型进行检测"
+    fn=predict_image,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
+        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
+    ],
+    outputs=gr.Image(type="pil", label="Result"),
+    title="Ultralytics Gradio",
+    description="Upload images for inference. The Ultralytics YOLOv8n model is used by default.",
+    examples=[
+        [ASSETS / "bus.jpg", 0.25, 0.45],
+        [ASSETS / "zidane.jpg", 0.25, 0.45],
+    ],
 )
 
-# 启动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
-)
+if __name__ == "__main__":
+    iface.launch()