Update predict.py

predict.py

@@ -1,39 +1,75 @@
-from PIL import Image
-from yolo import YOLO
-
-if __name__ == "__main__":
-    mode            = 'predict'
-    crop            = False
-    count           = True
-    dir_origin_path = "img/vs"
-    dir_save_path   = "img_out"
-  
-    yolo = YOLO()
-
-    if mode == "predict":
-        while True:
-            img = input('Input image filename:')
-            try:
-                image = Image.open(img)
-            except:
-                print('Open Error! Try again!')
-                continue
-            else:
-                r_image = yolo.detect_image(image, crop = crop, count=count)
-                r_image.show()
-
-    elif mode == "dir_predict":
-        import os
-        from tqdm import tqdm
-
-        img_names = os.listdir(dir_origin_path)
-        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(dir_origin_path, img_name)
-                image       = Image.open(image_path)
-                r_image     = yolo.detect_image(image)
-                if not os.path.exists(dir_save_path):
-                    os.makedirs(dir_save_path)
-                r_image.save(os.path.join(dir_save_path, img_name.replace(".jpg", ".png")), quality=95, subsampling=0)
-    else:
-        raise AssertionError("Please specify the correct mode: 'predict', 'dir_predict'.")
+from PIL import Image
+from yolo import YOLO
+import gradio as gr
+import os
+from tqdm import tqdm
+
+# Initialize YOLO model
+yolo = YOLO()
+
+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
+
+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 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.")
+
+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."
+
+css = """
+    .image-frame img, .image-container img {
+        width: auto;
+        height: auto;
+        max-width: none;
+    }
+"""
+
+demo = gr.Interface(
+    fn=inference,
+    inputs=[
+        gr.Image(type="pil", label="Input Image"),
+        gr.Radio(choices=["predict", "dir_predict"], label="Mode", value="predict"),
+        gr.Checkbox(value=False, label="Crop"),
+        gr.Checkbox(value=True, label="Count"),
+        gr.Textbox(placeholder="Input directory (for 'dir_predict' mode)", label="Input Directory", visible=False),
+        gr.Textbox(placeholder="Output directory (for 'dir_predict' mode)", label="Output Directory", visible=False),
+    ],
+    outputs=gr.Image(type="pil", label="Output Image"),
+    title=title,
+    description=description,
+    css=css,
+)
+
+if __name__ == "__main__":
+    demo.launch()