import os

from ultralytics import YOLO
import glob
# Load your trained segmentation model
model = YOLO('runs/segment/train/weights/last.pt')
IN_DIR = "/home/quentin/repos/yolov11-segmentation_earth-worm/data/yolo_split/images/val"
# Predict on an image, replace with your image path or URL

files = sorted(glob.glob(f"{IN_DIR}/*.jpg"))
for f in files:

    results = model(f)
# Save results, this includes saving masks
    os.makedirs("results/", exist_ok=True)
    results[0].save(f"results/result_{os.path.basename(f)}")
#
# for f in sorted(glob.glob(f"{IN_DIR}/2024-10-08/*.jpg")):
#     results = model(f)
#     # Save results, this includes saving masks
#     os.makedirs("results/2024-10-08", exist_ok=True)
#     results[0].save(f"results/2024-10-08/result_{os.path.basename(f)}")