Spaces:
Running
Running
| import os | |
| import subprocess | |
| # Clone the yolov5 repository and install its requirements | |
| if not os.path.exists('yolov5'): | |
| subprocess.run(['git', 'clone', 'https://github.com/ultralytics/yolov5'], check=True) | |
| subprocess.run(['pip', 'install', '-r', 'yolov5/requirements.txt'], check=True) | |
| import torch | |
| import torchvision | |
| from torchvision.transforms import functional as F | |
| from PIL import Image | |
| import cv2 | |
| import gradio as gr | |
| import numpy as np | |
| from yolov5.models.yolo import Model | |
| from yolov5.utils.general import non_max_suppression | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| print(f"Using device: {device}") | |
| model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True).to(device) | |
| model.eval() | |
| print("Model loaded successfully") | |
| def preprocess_image(image): | |
| try: | |
| 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 | |
| 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(image, cv2.COLOR_RGB2BGR) | |
| h, w, _ = image.shape | |
| for box in outputs: | |
| if box is not None: | |
| x1, y1, x2, y2, score, label = box[:6] | |
| if score > threshold: | |
| x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) | |
| cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2) | |
| text = f"{model.names[int(label)]:s}: {score:.2f}" | |
| cv2.putText(image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2) | |
| return cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| except Exception as e: | |
| print(f"Error in drawing boxes: {e}") | |
| return image | |
| def detect_objects(image): | |
| image_tensor = preprocess_image(image) | |
| if image_tensor is None: | |
| return image | |
| try: | |
| outputs = model(image_tensor)[0] # Get the first element of the output | |
| print(f"Model raw outputs: {outputs}") | |
| outputs = non_max_suppression(outputs, conf_thres=0.25, iou_thres=0.45)[0] # Apply NMS | |
| if outputs is None or len(outputs) == 0: | |
| print("No objects detected.") | |
| return image | |
| print(f"Filtered outputs: {outputs}") | |
| result_image = draw_boxes(image, outputs.cpu().numpy()) | |
| return result_image | |
| except Exception as e: | |
| print(f"Error in detecting objects: {e}") | |
| return image | |
| iface = gr.Interface( | |
| fn=detect_objects, | |
| 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." | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch() | |