Views
No views yet
1# Example Code: Try on google colab
2
3# Install required libraries
4!pip install ultralytics --quiet
5!pip install huggingface_hub --quiet
6import cv2
7import matplotlib.pyplot as plt
8from ultralytics import YOLO
9from huggingface_hub import hf_hub_download
10from google.colab import files
11import os
12
13# Download the YOLO model from Hugging Face
14model_path = hf_hub_download(repo_id="krishnamishra8848/Road_Detection", filename="best.pt")
15
16# Load the YOLO model
17model = YOLO(model_path)
18
19# Upload a photo
20print("Please upload an image:")
21uploaded = files.upload()
22
23for filename in uploaded.keys():
24 # Read the uploaded image
25 image = cv2.imread(filename)
26 image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
27
28 # Perform inference
29 results = model(image)
30
31 # Draw bounding boxes and class names
32 for result in results[0].boxes:
33 box = result.xyxy[0].cpu().numpy() # Bounding box (x_min, y_min, x_max, y_max)
34 cls = int(result.cls[0].cpu().numpy()) # Class ID
35 conf = result.conf[0].cpu().numpy() # Confidence score
36 label = f"{model.names[cls]}: {conf:.2f}" # Label with class name and confidence
37
38 # Draw the bounding box
39 cv2.rectangle(image_rgb, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 255, 0), 2)
40
41 # Draw the class name and confidence score
42 cv2.putText(image_rgb, label, (int(box[0]), int(box[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
43
44 # Display the image with bounding boxes
45 plt.figure(figsize=(10, 10))
46 plt.imshow(image_rgb)
47 plt.axis('off')
48 plt.show()
49
50 # Save the processed image
51 output_filename = "output_" + filename
52 cv2.imwrite(output_filename, cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR))
53 print(f"Processed image saved as {output_filename}")