Update app.py

app.py

@@ -7,6 +7,10 @@ from tqdm import tqdm
 # 初始化YOLO模型
 yolo = YOLO()
 
+# 获取 img 文件夹中的图片列表
+img_folder = "img"
+img_files = [os.path.join(img_folder, file) for file in os.listdir(img_folder) if file.endswith(('.png', '.jpg', '.jpeg'))]
+
 # 预测单张图像
 def predict_single_image(image, crop=False, count=True):
     try:
@@ -15,19 +19,22 @@ def predict_single_image(image, crop=False, count=True):
     except Exception as e:
         return str(e)
 
+# 创建图片选择器
+def select_image(image_path):
+    image = Image.open(image_path)
+    return image
 
 # 使用Gradio创建界面
-iface = gr.Interface(
-    fn=predict_single_image,
-    inputs=gr.inputs.Image(type="pil"),
-    outputs=gr.outputs.Image(type="pil"),
-    title="YOLO 图像检测器",
-    description="上传一张图像，使用YOLO模型进行对象检测。",
-    examples=[
-        ["img/1.png"],
-        ["img/2.png"]
-    ]
-)
+with gr.Blocks() as iface:
+    with gr.Row():
+        with gr.Column():
+            selected_image = gr.Image(label="Selected Image")
+            img_selector = gr.Dropdown(choices=img_files, label="Select an image from folder")
+            img_selector.change(select_image, img_selector, selected_image)
+        with gr.Column():
+            output_image = gr.Image(label="Output Image")
+            predict_button = gr.Button("Predict")
+            predict_button.click(predict_single_image, inputs=selected_image, outputs=output_image)
 
 # 启动应用
 if __name__ == "__main__":