from ultralytics import YOLO
import os
import torch
# Load model
# model = YOLO(f"/Users/duan_j/Applications/alc/tests/yolo/runs/detect/train{trainingset}/weights/best.pt")  # load pretrained model
#model = YOLO(f"models/best_v12_22092025.pt")  # load pretrained model
model_name = "best_yolo26n-seg-overlap-false_2026-04-12_nms"
model_type = "engine"
model = YOLO(f"/Users/duan_j/repos/aare_suite/AareLC/models/{model_name}.{model_type}")

prune = False
# Load the pruned state_dict
if prune:
   state_dict = torch.load(f"/Users/duan_j/repos/aare_suite/AareLC/models/pruned_{model_name}_weights.pt", map_location="cpu")
   model.model.load_state_dict(state_dict, strict=True)

#model.export(format="engine", conf=0.25, iou=0.45) this enables nms

# model = YOLO("yolov8n.yaml")  # build new model from YAML

# Export the model to TensorRT engine format
# model.export(format="engine", imgsz=640, simplify=False)  # Adjust imgsz if needed

# Continue training with new data
# results = model.train(
#     data="/Users/duan_j/Applications/alc/tests/yolo/dataset.yaml",
#     epochs=100,
# )
results = model.predict(
   # source = "/Users/duan_j/Applications/alc/tests/sample_img/20240109/only_images/A2",
   source="/Users/duan_j/repos/aare_suite/AareLC/testresults/testimgs",
   save=True,
   project=f"/Users/duan_j/repos/aare_suite/AareLC/testresults",
   name=f"{model_name}",
   conf=0.5  # Confidence threshold
)