Update app.py

app.py

@@ -1,61 +1,32 @@
 import gradio as gr
-import torch
-import torchvision.transforms as transforms
-from PIL import Image, ImageDraw
+from ultralytics import YOLO
+import cv2
+import numpy as np
+from PIL import Image
 
-# Load your YOLOv8 PyTorch model
+# Load your YOLO model
 model_path = 'best_hat_detection_model_100.pt'  # Your model file on Hugging Face Space
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = YOLO(model_path)
 
-# Load the model
-model = torch.load(model_path, map_location=device)['model'].float()
-model.eval()
-
-# Define transformations needed for your model
-transform = transforms.Compose([
-    transforms.Resize((640, 640)),
-    transforms.ToTensor(),
-])
-
-# Function to draw bounding boxes on the image
-def draw_bboxes(image, predictions, threshold=0.5):
-    draw = ImageDraw.Draw(image)
-    
-    for pred in predictions:
-        # YOLO model typically returns a list or tensor of bounding boxes and confidence scores
-        # pred: [x1, y1, x2, y2, confidence, class_id]
-        # Make sure to unpack these values correctly.
-        
-        x1, y1, x2, y2, confidence = pred[:4]  # Coordinates and confidence score
-        if confidence > threshold:  # Only draw boxes above the threshold
-            draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
-            draw.text((x1, y1), f"{confidence:.2f}", fill="red")
-    
-    return image
-
-# Function to predict and return image with bounding boxes
+# Function to perform object detection and return the image with bounding boxes
 def predict_safety_hat(image):
-    pil_image = Image.fromarray(image.astype('uint8'), 'RGB')  # Convert numpy array to PIL Image
-    input_tensor = transform(pil_image).unsqueeze(0).to(device)  # Apply transformations
-    
-    with torch.no_grad():
-        predictions = model(input_tensor)[0]  # Model output
+    # Convert image from numpy array (Gradio format) to a format YOLO can process
+    image = np.array(image)
     
-    # Assuming predictions is a tensor, convert it to a list
-    predictions = predictions.cpu().numpy()  # Convert tensor to numpy array
+    # Run inference using the YOLO model
+    results = model(image)
     
-    # Convert the image back to draw bounding boxes
-    pil_image = pil_image.convert("RGB")
+    # Get the image with predictions drawn (bounding boxes, etc.)
+    result_image = results[0].plot()  # YOLOv8 returns an image with the bounding boxes drawn
     
-    # Assuming `predictions` contains the bounding box coordinates and confidence scores
-    # Format of predictions: [x1, y1, x2, y2, confidence, class_id]
-    return draw_bboxes(pil_image, predictions)
+    # Convert the result image back to PIL format for Gradio
+    return Image.fromarray(result_image)
 
 # Gradio Interface
-inputs = gr.Image(label="Input Image")
-outputs = gr.Image(label="Output Image with Bounding Boxes")  # Change to output image
+inputs = gr.Image(label="Input Image", type="numpy")
+outputs = gr.Image(label="Output Image with Bounding Boxes")
 
-title = "Safety Hat Recognition"
-description = "Upload an image to detect safety hats and get bounding boxes."
+title = "Safety Hat Detection"
+description = "Upload an image to detect safety hats and see bounding boxes."
 
-gr.Interface(predict_safety_hat, inputs, outputs, title=title, description=description).launch()
+gr.Interface(fn=predict_safety_hat, inputs=inputs, outputs=outputs, title=title, description=description).launch()