36 lines
1.4 KiB
Python
36 lines
1.4 KiB
Python
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
|
|
) |