Update app.py

app.py

@@ -20,20 +20,7 @@ def save_image(image, filename):
 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
+    return result_image
 
 # Function to list images in the 'img' folder
 def get_image_list():
@@ -45,9 +32,8 @@ def get_image_list():
     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")
+image_output = gr.outputs.Image(type="pil", label="Output Image")
+image_select = gr.inputs.Radio(get_image_list, label="Select Image from 'img' Folder")
 
 # Gradio app for single image prediction
 iface = gr.Interface(
@@ -58,20 +44,29 @@ iface = gr.Interface(
     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")
+# Custom function to display images in a gallery format
+def gallery_update():
+    image_list = get_image_list()
+    return gr.Gallery.update(value=image_list)
 
-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."
-)
+# Gradio gallery component for displaying images in the 'img' folder
+gallery = gr.Gallery(get_image_list(), label="Image Gallery")
+
+# Update gallery when Gradio interface loads
+gallery.update(gallery_update)
 
 # Combine both interfaces
-app = gr.TabbedInterface([iface, iface_dir], ["Single Image Prediction", "Directory Prediction"])
+app = gr.Blocks()
+
+with app:
+    with gr.Row():
+        with gr.Column():
+            gallery.render()
+        with gr.Column():
+            image = gr.Image(type="pil", label="Selected Image")
+            predict_button = gr.Button("Predict")
+            output_image = gr.Image(type="pil", label="Output Image")
+            predict_button.click(predict, inputs=image, outputs=output_image)
 
 # Launch the app
 app.launch()