Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from PIL import Image, ImageDraw, ImageFont | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import os | |
| # Download images | |
| torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg', 'zidane.jpg') | |
| torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master/data/images/bus.jpg', 'bus.jpg') | |
| # Load YOLOv5 model | |
| model = torch.hub.load('ultralytics/yolov5', 'yolov5s') | |
| def yolo(im): | |
| try: | |
| # Check if the input is an Image object | |
| if isinstance(im, Image.Image): | |
| # Convert the PIL image to a numpy array | |
| im_array = np.array(im) | |
| # Perform inference with YOLOv5 | |
| results = model(im_array) # inference | |
| # Get the bounding boxes and labels | |
| boxes = results.xyxy[0].cpu().numpy() | |
| # Convert the results to a PIL Image | |
| output_image = Image.fromarray(im_array) | |
| # Draw the bounding boxes and labels on the output image | |
| draw = ImageDraw.Draw(output_image) | |
| font = ImageFont.load_default(45) | |
| for box in boxes: | |
| label = results.names[int(box[5])] | |
| draw.rectangle([(box[0], box[1]), (box[2], box[3])], outline="red", width=3) | |
| draw.text((box[0], box[1]), label, fill="blue", font=font) | |
| return output_image | |
| else: | |
| raise ValueError("The input should be an Image object.") | |
| except Exception as e: | |
| print(f"Error processing image: {e}") | |
| return None | |
| # Define Gradio interface | |
| inputs = gr.Image(type='pil', label="Original Image") | |
| outputs = gr.Image(type="pil", label="Output Image") | |
| title = "YOLOv5" | |
| description = "YOLOv5 Gradio demo for object detection. Upload an image or click an example image to use." | |
| article = "<p style='text-align: center'>YOLOv5 is a family of compound-scaled object detection models trained on the COCO dataset, and includes " \ | |
| "simple functionality for Test Time Augmentation (TTA), model ensembling, hyperparameter evolution, " \ | |
| "and export to ONNX, CoreML and TFLite. <a href='https://github.com/ultralytics/yolov5'>Source code</a> |" \ | |
| "<a href='https://apps.apple.com/app/id1452689527'>iOS App</a> | <a href='https://pytorch.org/hub/ultralytics_yolov5'>PyTorch Hub</a></p>" | |
| examples = [['zidane.jpg'], ['bus.jpg']] | |
| gr.Interface(yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch(debug=True) | |