Update app.py

app.py

@@ -1,67 +1,77 @@
-from PIL import Image
-from yolo import YOLO
-import gradio as gr
-import os
-
-# Initialize YOLO model
-yolo = YOLO()
-
-def detect_objects(image, crop=False, count=True):
-    r_image = yolo.detect_image(image, crop=crop, count=count)
-    return r_image
-
-def save_image(image, filename):
-    if not os.path.exists("img_out"):
-        os.makedirs("img_out")
-    image.save(os.path.join("img_out", filename), quality=95, subsampling=0)
-    return os.path.join("img_out", filename)
-
-# Gradio interface for single image prediction
-def predict(image):
-    result_image = detect_objects(image)
-    output_path = save_image(result_image, "output.png")
-    return output_path
-
-# Gradio interface for directory prediction
-def dir_predict(dir_origin_path):
-    img_names = os.listdir(dir_origin_path)
-    output_images = []
-    for img_name in img_names:
-        if img_name.lower().endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
-            image_path = os.path.join(dir_origin_path, img_name)
-            image = Image.open(image_path)
-            r_image = detect_objects(image)
-            output_path = save_image(r_image, img_name.replace(".jpg", ".png"))
-            output_images.append(output_path)
-    return output_images
-
-# Gradio interface components
-image_input = gr.inputs.Image(type="pil", label="Input Image")
-image_output = gr.outputs.Image(type="file", label="Output Image")
-
-# Gradio app
-iface = gr.Interface(
-    fn=predict, 
-    inputs=image_input, 
-    outputs=image_output,
-    title="YOLO Object Detection",
-    description="Upload an image to detect objects using YOLO model."
-)
-
-# Directory prediction interface
-dir_input = gr.inputs.Textbox(label="Directory Path")
-dir_output = gr.outputs.Textbox(label="Output Paths")
-
-iface_dir = gr.Interface(
-    fn=dir_predict,
-    inputs=dir_input,
-    outputs=dir_output,
-    title="YOLO Object Detection for Directory",
-    description="Provide a directory path to detect objects in all images within the directory."
-)
-
-# Combine both interfaces
-app = gr.TabbedInterface([iface, iface_dir], ["Single Image Prediction", "Directory Prediction"])
-
-# Launch the app
-app.launch()
+from PIL import Image
+from yolo import YOLO
+import gradio as gr
+import os
+
+# Initialize YOLO model
+yolo = YOLO()
+
+def detect_objects(image, crop=False, count=True):
+    r_image = yolo.detect_image(image, crop=crop, count=count)
+    return r_image
+
+def save_image(image, filename):
+    if not os.path.exists("img_out"):
+        os.makedirs("img_out")
+    image.save(os.path.join("img_out", filename), quality=95, subsampling=0)
+    return os.path.join("img_out", filename)
+
+# Gradio interface for single image prediction
+def predict(image):
+    result_image = detect_objects(image)
+    output_path = save_image(result_image, "output.png")
+    return output_path
+
+# Gradio interface for directory prediction
+def dir_predict(dir_origin_path):
+    img_names = os.listdir(dir_origin_path)
+    output_images = []
+    for img_name in img_names:
+        if img_name.lower().endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
+            image_path = os.path.join(dir_origin_path, img_name)
+            image = Image.open(image_path)
+            r_image = detect_objects(image)
+            output_path = save_image(r_image, img_name.replace(".jpg", ".png"))
+            output_images.append(output_path)
+    return output_images
+
+# Function to list images in the 'img' folder
+def get_image_list():
+    image_folder = "img"
+    if not os.path.exists(image_folder):
+        os.makedirs(image_folder)
+    img_names = [f for f in os.listdir(image_folder) if f.lower().endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff'))]
+    img_paths = [os.path.join(image_folder, img_name) for img_name in img_names]
+    return img_paths
+
+# Gradio interface components
+image_input = gr.inputs.Image(type="pil", label="Input Image")
+image_output = gr.outputs.Image(type="file", label="Output Image")
+image_select = gr.inputs.Dropdown(get_image_list, label="Select Image from 'img' Folder")
+
+# Gradio app for single image prediction
+iface = gr.Interface(
+    fn=predict, 
+    inputs=image_select, 
+    outputs=image_output,
+    title="YOLO Object Detection",
+    description="Select an image from the 'img' folder to detect objects using YOLO model."
+)
+
+# Directory prediction interface
+dir_input = gr.inputs.Textbox(default="img", label="Directory Path")
+dir_output = gr.outputs.Textbox(label="Output Paths")
+
+iface_dir = gr.Interface(
+    fn=dir_predict,
+    inputs=dir_input,
+    outputs=dir_output,
+    title="YOLO Object Detection for Directory",
+    description="Detect objects in all images within the 'img' directory."
+)
+
+# Combine both interfaces
+app = gr.TabbedInterface([iface, iface_dir], ["Single Image Prediction", "Directory Prediction"])
+
+# Launch the app
+app.launch()