|
|
--- |
|
|
license: apache-2.0 |
|
|
--- |
|
|
|
|
|
```python |
|
|
# Example Code: Try on google colab |
|
|
|
|
|
# Install required libraries |
|
|
!pip install ultralytics --quiet |
|
|
!pip install huggingface_hub --quiet |
|
|
import cv2 |
|
|
import matplotlib.pyplot as plt |
|
|
from ultralytics import YOLO |
|
|
from huggingface_hub import hf_hub_download |
|
|
from google.colab import files |
|
|
import os |
|
|
|
|
|
# Download the YOLO model from Hugging Face |
|
|
model_path = hf_hub_download(repo_id="krishnamishra8848/Road_Detection", filename="best.pt") |
|
|
|
|
|
# Load the YOLO model |
|
|
model = YOLO(model_path) |
|
|
|
|
|
# Upload a photo |
|
|
print("Please upload an image:") |
|
|
uploaded = files.upload() |
|
|
|
|
|
for filename in uploaded.keys(): |
|
|
# Read the uploaded image |
|
|
image = cv2.imread(filename) |
|
|
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) |
|
|
|
|
|
# Perform inference |
|
|
results = model(image) |
|
|
|
|
|
# Draw bounding boxes and class names |
|
|
for result in results[0].boxes: |
|
|
box = result.xyxy[0].cpu().numpy() # Bounding box (x_min, y_min, x_max, y_max) |
|
|
cls = int(result.cls[0].cpu().numpy()) # Class ID |
|
|
conf = result.conf[0].cpu().numpy() # Confidence score |
|
|
label = f"{model.names[cls]}: {conf:.2f}" # Label with class name and confidence |
|
|
|
|
|
# Draw the bounding box |
|
|
cv2.rectangle(image_rgb, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 255, 0), 2) |
|
|
|
|
|
# Draw the class name and confidence score |
|
|
cv2.putText(image_rgb, label, (int(box[0]), int(box[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) |
|
|
|
|
|
# Display the image with bounding boxes |
|
|
plt.figure(figsize=(10, 10)) |
|
|
plt.imshow(image_rgb) |
|
|
plt.axis('off') |
|
|
plt.show() |
|
|
|
|
|
# Save the processed image |
|
|
output_filename = "output_" + filename |
|
|
cv2.imwrite(output_filename, cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)) |
|
|
print(f"Processed image saved as {output_filename}") |
|
|
|