krishnamishra8848/Road_Detection

# 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}")