2023-DnB-Project/example/demo.py

from ultralytics import YOLO
# from ultralytics.yolo.utils.benchmarks import benchmark

# Load a model
# model = YOLO("yolov8n.yaml")  # build a new model from scratch
# model = YOLO("../runs/detect/train3/weights/best.pt")  # zhou
model = YOLO("../runs/detect/train4/weights/best.pt")  # 1000 img, mine
model = YOLO("../runs/detect/train2/weights/best.pt")  # 1000 img, based on 3000 img
# model = YOLO("../../../project/runs/detect/train3/weights/best.pt")  # 3000 img, mine


# Use the model
# model.train(data="coco128.yaml", epochs=3,workers=0)  # train the model,workers=0 if windows
# metrics = model.val()  # evaluate model performance on the validation set
img_path = "./image"
results = model.predict(img_path, save = True) # device=0 by default, conf:置信度阈值
# results = model.predict(img_path,save=True,classes=[0,2],conf=0.5) # i.e. classes=0,classes=[0,3,4]

# save detection results *
# results = model.predict(img_path,save=True,save_txt=True,classes=0,conf=0.4)