Update app.py

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 import os
-from helmet_detect_alert import detect_and_alert, detect_from_images, MODEL_PATHS
+import numpy as np
+from PIL import Image
+from helmet_detect_alert import detect_and_alert, MODEL_PATHS
 from ultralytics import YOLO
 
 # Load YOLOv8 model
@@ -8,14 +10,13 @@ model_path = MODEL_PATHS["YOLOv8n"]
 model = YOLO(model_path)
 
 # ---------- Detection Functions ----------
-
 def detect_image_fn(image, confidence):
     if image is None:
         return None
-    image_np = np.array(image)  # Convert PIL to NumPy
+    image_np = np.array(image)
     results = model.predict(image_np, conf=confidence, verbose=False)
     annotated = results[0].plot()
-    return annotated
+    return Image.fromarray(annotated)
 
 def detect_video_fn(video_file, confidence):
     if video_file is None: