Update app.py

app.py

@@ -1,88 +1,30 @@
-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(image, crop=False, count=True):
-    """
-    Predict single image using YOLO model
-    """
-    try:
-        print("Starting image prediction...")  # Debug log
-        r_image = yolo.detect_image(image, crop=crop, count=count)
-        print("Prediction completed.")  # Debug log
-        return r_image
-    except Exception as e:
-        print(f"Error during image prediction: {e}")  # Debug log
-        return None
+def predict(image):
+    r_image = yolo.detect_image(image)
+    return r_image
 
-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
-    """
-    try:
-        print(f"Processing directory: {input_dir}")  # Debug log
-        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))
-        print(f"Directory processing completed: {input_dir}")  # Debug log
-        return results
-    except Exception as e:
-        print(f"Error during directory prediction: {e}")  # Debug log
-        return []
-
-def inference(image, mode='predict', crop=False, count=True, input_dir=None, output_dir=None):
-    try:
-        print(f"Received mode: {mode}, crop: {crop}, count: {count}")  # Debug log
-        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.")
-    except Exception as e:
-        print(f"Error in inference function: {e}")  # Debug log
-        return None
-
-# 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",
-    "img/2.png",
-    "img/3.png",
-    "img/4.png",
-    "img/5.png",
-    "img/6.png",
-    "img/7.png",
-    "img/8.png",
-    "img/10.png",
-]
+title = "MASFNet: Multi-scale Adaptive Sampling Fusion Network for Object Detection in Adverse Weather "
+description = ""
+article = ""
 
 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;
-    }
-"""
+example_images = [
+    ["img/1.png"],
+    ["img/2.png"],
+    ["img/3.png"],
+    ["img/4.png"],
+    ["img/5.png"],
+    ["img/6.png"],
+    ["img/7.png"],
+    ["img/8.png"],
+    ["img/9.png"],
+]
 
 with gr.Blocks() as demo:
     gr.Markdown(f"### {title}")
@@ -91,20 +33,16 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             img_input = gr.Image(type="pil", label="Upload an Image")
-            mode_input = gr.Radio(choices=["predict", "dir_predict"], label="Mode", value="predict")
-            crop_input = gr.Checkbox(value=False, label="Crop")
-            count_input = gr.Checkbox(value=True, label="Count")
             submit_btn = gr.Button("Submit")
         with gr.Column():
             output = gr.Image(type="pil", label="Prediction Result")
     
-    submit_btn.click(fn=inference, inputs=[img_input, mode_input, crop_input, count_input], outputs=output)
+    submit_btn.click(fn=predict, inputs=img_input, outputs=output)
     demo.load(reset_interface, None, [output])
-
-    with gr.Row():
-        gr.Examples(
-            examples=example_images,
-            inputs=img_input,
-        )
+    gr.Examples(
+        examples=example_images,
+        inputs=img_input,
+    )
+    
 
 demo.launch()