from ultralytics import YOLO
import cv2

# 1. Load your trained model
# Adjust the path if your run was 'train2', 'train3', etc.
model = YOLO("runs/segment/train/weights/best.pt")

# 2. Run inference on a new image
# source can be an image path, video path, or '0' for webcam
results = model.predict(source="example.bmp", conf=0.5, save=True)

# 3. Show results
for r in results:
    # Plot the result (draws masks on the image)
    im_array = r.plot()  
    
    # Display using OpenCV
    cv2.imshow("YOLO11 Prediction", im_array)
    cv2.waitKey(0) # Wait for a key press to close

cv2.destroyAllWindows()