Update app.py

app.py

@@ -12,7 +12,7 @@ from torchvision.transforms import functional as F
 from PIL import Image
 import cv2
 import gradio as gr
-import numpy as np  # Add this import
+import numpy as np
 
 from yolov5.models.yolo import Model
 from yolov5.utils.general import non_max_suppression
@@ -25,17 +25,17 @@ print("Model loaded successfully")
 
 def preprocess_image(image):
     try:
-        image = image.convert("RGB")  # Ensure image is in RGB mode
-        image_tensor = F.to_tensor(image)
+        image = Image.fromarray(image)  # Convert numpy array to PIL Image
+        image_tensor = F.to_tensor(image).unsqueeze(0).to(device)
         print(f"Preprocessed image tensor: {image_tensor.shape}")
-        return image_tensor.unsqueeze(0).to(device)
+        return image_tensor
     except Exception as e:
         print(f"Error in preprocessing image: {e}")
         return None
 
 def draw_boxes(image, outputs, threshold=0.3):
     try:
-        image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
         h, w, _ = image.shape
 
         for box in outputs:
@@ -64,7 +64,7 @@ def detect_objects(image):
             print("No objects detected.")
             return image
         print(f"Filtered outputs: {outputs}")
-        result_image = draw_boxes(image, outputs)
+        result_image = draw_boxes(image, outputs.cpu().numpy())
         return result_image
     except Exception as e:
         print(f"Error in detecting objects: {e}")
@@ -72,9 +72,11 @@ def detect_objects(image):
 
 iface = gr.Interface(
     fn=detect_objects,
-    inputs=gr.Image(type="pil"),  # Removed source="webcam"
-    outputs=gr.Image(type="pil"),
-    live=True
+    inputs=gr.Image(type="numpy"),
+    outputs=gr.Image(type="numpy"),
+    title="YOLOv5 Object Detection",
+    description="Upload an image to detect objects using the YOLOv5 model."
 )
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()