Update app.py

app.py

@@ -12,37 +12,50 @@ 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: {e}")
+        print(f"Error during image prediction: {e}")  # Debug log
         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
+    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):
-    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.")
+    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"
@@ -95,4 +108,3 @@ with gr.Blocks() as demo:
         )
 
 demo.launch()
-