For this, you will need to have a .pt file for your YOLO model. This will NOT work with an OBB model, you will need to make some changes to the code.
Notes:
• Output Path: Update the output path for saved images as needed.
• Environment: This tutorial/code works on systems with OpenCV, a working camera, and a Python environment with the required libraries installed.
• YOLO Model: Ensure model.pt or your specific YOLO model file is correctly placed and compatible with the ultralytics YOLO library.
Step 1: Import Required Libraries
Ensure you have the necessary libraries installed. The key libraries are OpenCV, YOLO (from ultralytics), and PIL for handling images.
import cv2
from ultralytics import YOLO
Step 2: Load the YOLO Model
Specify the path to your YOLO model file and load it.
# Load the YOLO model
model = YOLO("model.pt") # Replace with your model file
Step 3: Set Up the Camera
Initialise the camera using OpenCV to capture video frames.
cap = cv2.VideoCapture(0) # 0 for default camera
if not cap.isOpened():
print("Error: Cannot access the camera.")
exit()
Step 4: Process Frames and Run YOLO Detection
Read frames from the camera, resize them for YOLO, and pass them to the model for detection.
while True:
ret, frame = cap.read()
if not ret:
print("Error: Cannot read frame.")
break
# Resize the frame for YOLO
frame_resized = cv2.resize(frame, (640, 480))
# Run YOLO detection
results = model(frame_resized)
# Extract detections
for box in results[0].boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0]) # Bounding box coordinates
label = f"{model.names[int(box.cls)]} ({box.conf:.2f})" # Label with confidence
# Draw the bounding box
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) # Green box
# Draw the label
cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
Step 5: Save the Annotated Frames
Save the processed frame with annotations to a file.
# Save the annotated frame
output_path = "annotated_frame.jpg"
cv2.imwrite(output_path, frame)
print(f"Saved annotated frame to {output_path}")
Step 6: Display the Annotated Frames (Optional)
Show the annotated frames in a window to verify the results in real-time.
# Display the annotated frame
cv2.imshow("YOLO Detection", frame)
# Exit loop when 'q' is pressed
if cv2.waitKey(1) & 0xFF == ord('q'):
break
Step 7: Clean Up Resources
Release the camera and close all OpenCV windows when the program finishes or exits.
cap.release()
cv2.destroyAllWindows()
Full Code Example
Here is the complete code integrating all the steps:
import cv2
from ultralytics import YOLO
#Load the YOLO model
model = YOLO("model.pt") # Replace with your model file
#Set up the camera
cap = cv2.VideoCapture(0) # 0 for default camera
if not cap.isOpened():
print("Error: Cannot access the camera.")
exit()
while True:
ret, frame = cap.read()
if not ret:
print("Error: Cannot read frame.")
break
# Resize the frame for YOLO
frame_resized = cv2.resize(frame, (640, 480))
# Run YOLO detection
results = model(frame_resized)
# Extract detections
for box in results[0].boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0]) # Bounding box coordinates
label = f"{model.names[int(box.cls)]} ({box.conf:.2f})" # Label with confidence
# Draw the bounding box
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) # Green box
# Draw the label
cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
# Save the annotated frame
output_path = "annotated_frame.jpg"
cv2.imwrite(output_path, frame)
print(f"Saved annotated frame to {output_path}")
# Display the annotated frame
cv2.imshow("YOLO Detection", frame)
# Exit loop when 'q' is pressed
if cv2.waitKey(1) & 0xFF == ord('q'):
break
#Release resources
cap.release()
cv2.destroyAllWindows()