from ultralytics import YOLO

# Load a model
model = YOLO("best.pt")  # load a pretrained model (recommended for training)

# Use the model
# model.train(data="/data.yaml", epochs=3)  # train the model
# metrics = model.val()  # evaluate model performance on the validation set
results = model("data/Screenshot 2023-10-21 at 14.28.32.png", stream=True, show=True, save=True, conf=0.7)  # predict on an image
# path = model.export(format="onnx")  # export the model to ONNX format

# Process results generator
for result in results:
    boxes = result.boxes  # Boxes object for bbox outputs
    masks = result.masks  # Masks object for segmentation masks outputs
    keypoints = result.keypoints  # Keypoints object for pose outputs
    probs = result.probs  # Probs object for classification outputs