2023-DnB-Project/example/demo.py
2023-06-29 23:49:53 +08:00

21 lines
1007 B
Python

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)