File size: 1,299 Bytes
605bd6c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
import cv2
import matplotlib.pyplot as plt
from ultralytics import YOLO
# Define a class for YOLO-based prediction
class YOLOPredictor:
def __init__(self, model_path):
# Initialize the YOLO model with the given model path
self.model = YOLO(model_path)
def predict(self, image_path):
# Use the model to predict objects in the image at the given path
results = self.model.predict(image_path)
# Plot the results on the image
result_img = results[0].plot()
# Convert the image from BGR to RGB (OpenCV uses BGR by default)
result_img = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
# Display the result image using matplotlib
plt.imshow(result_img)
plt.show()
# Check if the script is run directly (not imported as a module)
if __name__ == "__main__":
# Define the path to the YOLO model
model_path = 'model/best.pt'
# Define the path to the image to be processed
image_path = 'dataset/test/images/20230607-012140_jpg.rf.8c726235ab75ba1861f667676087e1dd.jpg'
# Create a YOLOPredictor object with the specified model path
predictor = YOLOPredictor(model_path)
# Predict and display the result for the given image
predictor.predict(image_path)
|