Update app.py

app.py

@@ -1,17 +1,52 @@
+import os
+from PIL import Image
 import gradio as gr
+from tqdm import tqdm
 from yolo import YOLO
 
+# Initialize YOLO model
 yolo = YOLO()
 
-def predict(image):
-    r_image = yolo.detect_image(image)
-    return r_image
+def predict_image(image, crop=False, count=True):
+    """
+    Predict single image using YOLO model
+    """
+    try:
+        r_image = yolo.detect_image(image, crop=crop, count=count)
+        return r_image
+    except Exception as e:
+        print(f"Error: {e}")
+        return None
 
-title = "MASFNet: Multiscale Adaptive Sampling Fusion Network for Object Detection in Adverse Weather"
-description = ""
+def predict_directory(input_dir, output_dir, crop=False, count=True):
+    """
+    Predict images in a directory using YOLO model and save results to another directory
+    """
+    img_names = os.listdir(input_dir)
+    results = []
+    for img_name in tqdm(img_names):
+        if img_name.lower().endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
+            image_path = os.path.join(input_dir, img_name)
+            image = Image.open(image_path)
+            r_image = yolo.detect_image(image, crop=crop, count=count)
+            if not os.path.exists(output_dir):
+                os.makedirs(output_dir)
+            output_path = os.path.join(output_dir, img_name.replace(".jpg", ".png"))
+            r_image.save(output_path, quality=95, subsampling=0)
+            results.append((img_name, output_path))
+    return results
 
-def reset_interface():
-    return gr.update(value=None), gr.update(visible=False)
+def inference(image, mode='predict', crop=False, count=True, input_dir=None, output_dir=None):
+    if mode == 'predict':
+        return predict_image(image, crop=crop, count=count)
+    elif mode == 'dir_predict' and input_dir and output_dir:
+        return predict_directory(input_dir, output_dir, crop=crop, count=count)
+    else:
+        raise ValueError("Invalid mode or missing directories for 'dir_predict' mode.")
+
+# Gradio interface setup
+title = "YOLO Image Prediction"
+description = "This demo allows you to perform image prediction using a YOLO model. You can either predict a single image or all images in a directory."
 
 example_images = [
     "img/1.png",
@@ -25,6 +60,17 @@ example_images = [
     "img/10.png",
 ]
 
+def reset_interface():
+    return gr.update(value=None), gr.update(visible=False)
+
+css = """
+    .image-frame img, .image-container img {
+        width: auto;
+        height: auto;
+        max-width: none;
+    }
+"""
+
 with gr.Blocks() as demo:
     gr.Markdown(f"### {title}")
     gr.Markdown(description)
@@ -36,7 +82,7 @@ with gr.Blocks() as demo:
         with gr.Column():
             output = gr.Image(type="pil", label="Prediction Result")
     
-    submit_btn.click(fn=predict, inputs=img_input, outputs=output)
+    submit_btn.click(fn=inference, inputs=[img_input, 'predict', False, True], outputs=output)
     demo.load(reset_interface, None, [output])
 
     with gr.Row():