from ultralytics import YOLO
from pathlib import Path


def foo():
    model = YOLO("yolo11x.pt")
    example = Path(
        '/home/didier/all-dressed-programming/dodgeball-tracker/data/good-images-for-training/PB120196.JPG')
    results = model.predict(example)
    print('print done predicting')
    print(f'done predicting; {len(results)} results')
    for res in results:
        print(f"    >>> {res.boxes.data}")

def main():
    # Load a model
    model = YOLO("yolo11s.pt")  # load a pretrained model (recommended for training)

    # Train the model
    train = '/home/didier/all-dressed-programming/dodgeball-tracker/data/wheels_004.yaml'
    model.train(data=Path(train),
                epochs=150,
                imgsz=640,
                cache='ram')
    print('Done training')
    bad_example = Path("/home/didier/all-dressed-programming/dodgeball-tracker/data/label-studio-exports/2024-11-12_001/images/train/2a6267f7-PB120166.JPG")
    example = Path('/home/didier/all-dressed-programming/dodgeball-tracker/data/good-images-for-training/PB120196.JPG')
    results = model.predict(bad_example)
    print(f'done predicting; {len(results)} results')
    for res in results:
        print(f"    >>> {res.boxes.data}")

if __name__ == '__main__':
    foo()