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)