导出onnx，int8量化，评估，推理
pt转onnx
from ultralytics import YOLO
model = YOLO("./runs/segment/train6/weights/best.pt") # load a pretrained YOLOv8n model
model.export(format="onnx",opset=12) # export the model to ONNX forma
这里我们直接导出的话，默认的batch数量是1，也就是13640*640的输入。如果需要进行多batch的onnx导出，需要去ultralytics/yolo/engine/model.py里进行修改，第269行，将1改为你想要的batch数量：
self._check_is_pytorch_model()
overrides = self.overrides.copy()
overrides.update(kwargs)
args = get_cfg(cfg=DEFAULT_CFG, overrides=overrides)
args.task = self.task
if args.imgsz == DEFAULT_CFG.imgsz:
args.imgsz = self.model.args['imgsz'] # use trained imgsz unless custom value is passed
if args.batch == DEFAULT_CFG.batch:
args.batch = 10 # default to 1 if not modified
exporter = Exporter(overrides=args)
return exporter(model=self.model)
这样再导出的话，我们可以通过netron来查看：
此时可以发现，我们的batch已经从1变成了10了。目前我在我的openvino专栏中已经发布了基于batch=1时的推理，后续会跟进在c++中的openvino的batch推理。
1. 量化
import nncf
from openvino.tools import mo
from openvino.runtime import serialize
import torch
from pathlib import Path
import logging
import os
from zipfile import ZipFile
from multiprocessing.pool import ThreadPool
import yaml
from itertools import repeat
import time
import platform
from tqdm.notebook import tqdm
from ultralytics.yolo.utils.metrics import ConfusionMatrix
import torch
import numpy as np
def my_test(model, core, data_loader:torch.utils.data.DataLoader, validator, num_samples:int = None):
"""
OpenVINO YOLOv8 model accuracy validation function. Runs model validation on dataset and returns metrics
Parameters:
model (Model): OpenVINO model
data_loader (torch.utils.data.DataLoader): dataset loader
validato: instalce of validator class
num_samples (int, *optional*, None): validate model only on specified number samples, if provided
Returns:
stats: (Dict[str, float]) - dictionary with aggregated accuracy metrics statistics, key is metric name, value is metric value
"""
validator.seen = 0
validator.jdict = []
validator.stats = []
validator.batch_i = 1
validator.confusion_matrix = ConfusionMatrix(nc=validator.nc)
model.reshape({0: [1, 3, -1, -1]})
num_outputs = len(model.outputs)
compiled_model = core.compile_model(model)
for batch_i, batch in enumerate(tqdm(data_loader, total=num_samples)):
if num_samples is not None and batch_i == num_samples:
break
batch = validator.preprocess(batch)
results = compiled_model(batch["img"])
if num_outputs == 1:
preds = torch.from_numpy(results[compiled_model.output(0)])
else:
preds = [torch.from_numpy(results[compiled_model.output(0)]), torch.from_numpy(results[compiled_model.output(1)])]
preds = validator.postprocess(preds)
validator.update_metrics(preds, batch)
stats = validator.get_stats()
return stats
def print_stats(stats:np.ndarray, total_images:int, total_objects:int):
"""
Helper function for printing accuracy statistic
Parameters:
stats: (Dict[str, float]) - dictionary with aggregated accuracy metrics statistics, key is metric name, value is metric value
total_images (int) - number of evaluated images
total objects (int)
Returns:
None
"""
# print("Boxes:")
mp, mr, map50, mean_ap = stats['metrics/precision(B)'], stats['metrics/recall(B)'], stats['metrics/mAP50(B)'], stats['metrics/mAP50-95(B)']
# Print results
s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Labels', 'Precision', 'Recall', 'mAP@.5', 'mAP@.5:.95')
print('Boxes:',s)
pf = '%20s' + '%12i' * 2 + '%12.3g' * 4 # print format
print(pf % ('\t\t\tall', total_images, total_objects, mp, mr, map50, mean_ap))
if 'metrics/precision(M)' in stats:
# print("Masks:")
s_mp, s_mr, s_map50, s_mean_ap = stats['metrics/precision(M)'], stats['metrics/recall(M)'], stats['metrics/mAP50(M)'], stats['metrics/mAP50-95(M)']
# Print results
s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Labels', 'Precision', 'Recall', 'mAP@.5', 'mAP@.5:.95')
print("Masks:",s)
pf = '%20s' + '%12i' * 2 + '%12.3g' * 4 # print format
print(pf % ('\t\t\tall', total_images, total_objects, s_mp, s_mr, s_map50, s_mean_ap))
if __name__ == "__main__":
#定义需要量化的onnx模型的路径以及量化后保存的名称
MODEL_NAME = "best.onnx"
MODEL_PATH = f"runs/segment/train6/weights"
IR_MODEL_NAME = "v8_seg"
#将需要量化的onnx模型路径串起来
onnx_path = f"{MODEL_PATH}/{MODEL_NAME}"
#定义量化后的模型的路径名称
FP32_path = f"{MODEL_PATH}/FP32_openvino_model/{IR_MODEL_NAME}_FP32.xml"
FP16_path = f"{MODEL_PATH}/FP16_openvino_model/{IR_MODEL_NAME}_FP16.xml"
Int8_path = f"{MODEL_PATH}/Int8_openvino_model/{IR_MODEL_NAME}_Int8.xml"
#FP32 model
model = mo.convert_model(onnx_path)
serialize(model, FP32_path)
print(f"export ONNX to Openvino FP32 IR to:{FP32_path}")
#FP16 model
model = mo.convert_model(onnx_path, compress_to_fp16=True)
serialize(model, FP16_path)
print(f"export ONNX to Openvino FP16 IR to:{FP16_path}")
#Int8 model
from ultralytics.yolo.data.utils import check_det_dataset
from ultralytics.yolo.data.dataloaders.v5loader import create_dataloader
#
def create_data_source():
data = check_det_dataset('ultralytics/datasets/coco128-seg.yaml')
val_dataloader = create_dataloader(data['val'], imgsz=640, batch_size=1, stride=32, pad=0.5, workers=1)[0]
return val_dataloader
def transform_fn(data_item):
images = data_item['img']
images = images.float()
images = images / 255.0
images = images.cpu().detach().numpy()
return images
#加载数据
data_source = create_data_source()
#实例化校准数据集
nncf_calibration_dataset = nncf.Dataset(data_source, transform_fn)
#配置量化管道
subset_size = 40
preset = nncf.QuantizationPreset.MIXED
#执行模型量化
from openvino.runtime import Core
from openvino.runtime import serialize
core = Core()
ov_model = core.read_model(FP16_path)
quantized_model = nncf.quantize(
ov_model, nncf_calibration_dataset, preset=preset, subset_size=subset_size
)
serialize(quantized_model, Int8_path)
print(f"export ONNX to Openvino Int8 IR to:{Int8_path}")
2. 评估
import nncf
from openvino.tools import mo
from openvino.runtime import serialize
import torch
from pathlib import Path
import logging
import os
from zipfile import ZipFile
from multiprocessing.pool import ThreadPool
import yaml
from itertools import repeat
import time
import platform
from tqdm.notebook import tqdm
from ultralytics.yolo.utils.metrics import ConfusionMatrix
import torch
import numpy as np
def my_test(model, core, data_loader:torch.utils.data.DataLoader, validator, num_samples:int = None):
"""
OpenVINO YOLOv8 model accuracy validation function. Runs model validation on dataset and returns metrics
Parameters:
model (Model): OpenVINO model
data_loader (torch.utils.data.DataLoader): dataset loader
validato: instalce of validator class
num_samples (int, *optional*, None): validate model only on specified number samples, if provided
Returns:
stats: (Dict[str, float]) - dictionary with aggregated accuracy metrics statistics, key is metric name, value is metric value
"""
validator.seen = 0
validator.jdict = []
validator.stats = []
validator.batch_i = 1
validator.confusion_matrix = ConfusionMatrix(nc=validator.nc)
model.reshape({0: [1, 3, -1, -1]})
num_outputs = len(model.outputs)
compiled_model = core.compile_model(model)
for batch_i, batch in enumerate(tqdm(data_loader, total=num_samples)):
if num_samples is not None and batch_i == num_samples:
break
batch = validator.preprocess(batch)
results = compiled_model(batch["img"])
if num_outputs == 1:
preds = torch.from_numpy(results[compiled_model.output(0)])
else:
preds = [torch.from_numpy(results[compiled_model.output(0)]), torch.from_numpy(results[compiled_model.output(1)])]
preds = validator.postprocess(preds)
validator.update_metrics(preds, batch)
stats = validator.get_stats()
return stats
def print_stats(stats:np.ndarray, total_images:int, total_objects:int):
"""
Helper function for printing accuracy statistic
Parameters:
stats: (Dict[str, float]) - dictionary with aggregated accuracy metrics statistics, key is metric name, value is metric value
total_images (int) - number of evaluated images
total objects (int)
Returns:
None
"""
# print("Boxes:")
mp, mr, map50, mean_ap = stats['metrics/precision(B)'], stats['metrics/recall(B)'], stats['metrics/mAP50(B)'], stats['metrics/mAP50-95(B)']
# Print results
s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Labels', 'Precision', 'Recall', 'mAP@.5', 'mAP@.5:.95')
print('Boxes:',s)
pf = '%20s' + '%12i' * 2 + '%12.3g' * 4 # print format
print(pf % ('\t\t\tall', total_images, total_objects, mp, mr, map50, mean_ap))
if 'metrics/precision(M)' in stats:
# print("Masks:")
s_mp, s_mr, s_map50, s_mean_ap = stats['metrics/precision(M)'], stats['metrics/recall(M)'], stats['metrics/mAP50(M)'], stats['metrics/mAP50-95(M)']
# Print results
s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Labels', 'Precision', 'Recall', 'mAP@.5', 'mAP@.5:.95')
print("Masks:",s)
pf = '%20s' + '%12i' * 2 + '%12.3g' * 4 # print format
print(pf % ('\t\t\tall', total_images, total_objects, s_mp, s_mr, s_map50, s_mean_ap))
if __name__ == "__main__":
from ultralytics.yolo.utils import DEFAULT_CFG
from ultralytics.yolo.cfg import get_cfg
from ultralytics.yolo.data.utils import check_det_dataset
from ultralytics import YOLO
from ultralytics.yolo.utils import ops
from openvino.runtime import Core
MODEL_PATH = f"runs/segment/train6/weights"
IR_MODEL_NAME = "v8_seg"
core = Core()
pt_path = f"{MODEL_PATH}/best.pt"
#将需要量化的onnx模型路径串起来
# onnx_path = f"{MODEL_PATH}/{MODEL_NAME}"
#定义量化后的模型的路径名称
FP32_path = f"{MODEL_PATH}/FP32_openvino_model/{IR_MODEL_NAME}_FP32.xml"
FP16_path = f"{MODEL_PATH}/FP16_openvino_model/{IR_MODEL_NAME}_FP16.xml"
Int8_path = f"{MODEL_PATH}/Int8_openvino_model/{IR_MODEL_NAME}_Int8.xml"
CFG_PATH = 'ultralytics/yolo/cfg/default.yaml'
NUM_TEST_SAMPLES = 300
args = get_cfg(cfg=DEFAULT_CFG)
args.data = str(CFG_PATH)
fp32_model = core.read_model(FP32_path)
fp16_model = core.read_model(FP16_path)
quantized_model = core.read_model(Int8_path)
seg_model = YOLO(pt_path)
seg_validator = seg_model.ValidatorClass(args=args)
seg_validator.data = check_det_dataset('ultralytics/datasets/coco128-seg.yaml')
seg_data_loader = seg_validator.get_dataloader("D:\\Ultralytics\\datasets\\coco128-seg", 1)
seg_validator.is_coco = True
seg_validator.class_map = [1,2,3,4,5]
seg_validator.names = seg_model.model.names
seg_validator.metrics.names = seg_validator.names
seg_validator.nc = seg_model.model.model[-1].nc
seg_validator.nm = 32
seg_validator.process = ops.process_mask
seg_validator.plot_masks = []
# seg_data_loader = create_data_source()
fp32_seg_stats = my_test(fp32_model, core, seg_data_loader, seg_validator, num_samples=None)
fp16_seg_stats = my_test(fp16_model, core, seg_data_loader, seg_validator, num_samples=None)
int8_seg_stats = my_test(quantized_model, core, seg_data_loader, seg_validator, num_samples=None)
print("FP32 model accuracy")
print_stats(fp32_seg_stats, seg_validator.seen, seg_validator.nt_per_class.sum())
print("FP16 model accuracy")
print_stats(fp16_seg_stats, seg_validator.seen, seg_validator.nt_per_class.sum())
print("INT8 model accuracy")
print_stats(int8_seg_stats, seg_validator.seen, seg_validator.nt_per_class.sum())
3. python中的openvino推理
from openvino.runtime import Core, Model
from typing import Tuple
from ultralytics.yolo.utils import ops
import torch
import numpy as np
import cv2
from PIL import Image
from typing import Tuple, Dict
import cv2
from PIL import Image
from ultralytics.yolo.utils.plotting import colors
import time
def plot_one_box(box: np.ndarray, img: np.ndarray, color: Tuple[int, int, int] = None, mask: np.ndarray = None,
label: str = None, line_thickness: int = 5):
"""
Helper function for drawing single bounding box on image
Parameters:
x (np.ndarray): bounding box coordinates in format [x1, y1, x2, y2]
img (no.ndarray): input image
color (Tuple[int, int, int], *optional*, None): color in BGR format for drawing box, if not specified will be selected randomly
mask (np.ndarray, *optional*, None): instance segmentation mask polygon in format [N, 2], where N - number of points in contour, if not provided, only box will be drawn
label (str, *optonal*, None): box label string, if not provided will not be provided as drowing result
line_thickness (int, *optional*, 5): thickness for box drawing lines
"""
# Plots one bounding box on image img
tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness
color = color or [np.random.randint(0, 255) for _ in range(3)]
c1, c2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
if label:
tf = max(tl - 1, 1) # font thickness
t_size = cv2.getTextSize(label, 0, fontScale=tl, thickness=tf)[0]
c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled
cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl, [225, 255, 255], thickness=tf * 2, lineType=cv2.LINE_AA)
if mask is not None:
image_with_mask = img.copy()
mask
cv2.fillPoly(image_with_mask, pts=[mask.astype(int)], color=color)
img = cv2.addWeighted(img, 0.5, image_with_mask, 0.5, 1)
return img
def draw_results(results: Dict, source_image: np.ndarray, label_map: Dict):
"""
Helper function for drawing bounding boxes on image
Parameters:
image_res (np.ndarray): detection predictions in format [x1, y1, x2, y2, score, label_id]
source_image (np.ndarray): input image for drawing
label_map; (Dict[int, str]): label_id to class name mapping
Returns:
"""
boxes = results["det"]
masks = results.get("segment")
h, w = source_image.shape[:2]
for idx, (*xyxy, conf, lbl) in enumerate(boxes):
label = f'{label_map[int(lbl)]} {conf:.2f}'
mask = masks[idx] if masks is not None else None
source_image = plot_one_box(xyxy, source_image, mask=mask, label=label, color=colors(int(lbl)),
line_thickness=1)
return source_image
def letterbox(img: np.ndarray, new_shape: Tuple[int, int] = (640, 640), color: Tuple[int, int, int] = (114, 114, 114),
auto: bool = False, scale_fill: bool = False, scaleup: bool = False, stride: int = 32):
"""
Resize image and padding for detection. Takes image as input,
resizes image to fit into new shape with saving original aspect ratio and pads it to meet stride-multiple constraints
Parameters:
img (np.ndarray): image for preprocessing
new_shape (Tuple(int, int)): image size after preprocessing in format [height, width]
color (Tuple(int, int, int)): color for filling padded area
auto (bool): use dynamic input size, only padding for stride constrins applied
scale_fill (bool): scale image to fill new_shape
scaleup (bool): allow scale image if it is lower then desired input size, can affect model accuracy
stride (int): input padding stride
Returns:
img (np.ndarray): image after preprocessing
ratio (Tuple(float, float)): hight and width scaling ratio
padding_size (Tuple(int, int)): height and width padding size
"""
# Resize and pad image while meeting stride-multiple constraints
shape = img.shape[:2] # current shape [height, width]
if isinstance(new_shape, int):
new_shape = (new_shape, new_shape)
# Scale ratio (new / old)
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
if not scaleup: # only scale down, do not scale up (for better test mAP)
r = min(r, 1.0)
# Compute padding
ratio = r, r # width, height ratios
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
if auto: # minimum rectangle
dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding
elif scale_fill: # stretch
dw, dh = 0.0, 0.0
new_unpad = (new_shape[1], new_shape[0])
ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios
dw /= 2 # divide padding into 2 sides
dh /= 2
if shape[::-1] != new_unpad: # resize
img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
return img, ratio, (dw, dh)
def preprocess_image(img0: np.ndarray):
"""
Preprocess image according to YOLOv8 input requirements.
Takes image in np.array format, resizes it to specific size using letterbox resize and changes data layout from HWC to CHW.
Parameters:
img0 (np.ndarray): image for preprocessing
Returns:
img (np.ndarray): image after preprocessing
"""
# resize
img = letterbox(img0)[0]
# Convert HWC to CHW
img = img.transpose(2, 0, 1)
img = np.ascontiguousarray(img)
return img
def image_to_tensor(image: np.ndarray):
"""
Preprocess image according to YOLOv8 input requirements.
Takes image in np.array format, resizes it to specific size using letterbox resize and changes data layout from HWC to CHW.
Parameters:
img (np.ndarray): image for preprocessing
Returns:
input_tensor (np.ndarray): input tensor in NCHW format with float32 values in [0, 1] range
"""
input_tensor = image.astype(np.float32) # uint8 to fp32
input_tensor /= 255.0 # 0 - 255 to 0.0 - 1.0
# add batch dimension
if input_tensor.ndim == 3:
input_tensor = np.expand_dims(input_tensor, 0)
return input_tensor
def postprocess(
pred_boxes:np.ndarray,
input_hw:Tuple[int, int],
orig_img:np.ndarray,
min_conf_threshold:float = 0.25,
nms_iou_threshold:float = 0.7,
agnosting_nms:bool = False,
max_detections:int = 300,
pred_masks:np.ndarray = None,
retina_mask:bool = False
):
"""
YOLOv8 model postprocessing function. Applied non maximum supression algorithm to detections and rescale boxes to original image size
Parameters:
pred_boxes (np.ndarray): model output prediction boxes
input_hw (np.ndarray): preprocessed image
orig_image (np.ndarray): image before preprocessing
min_conf_threshold (float, *optional*, 0.25): minimal accepted confidence for object filtering
nms_iou_threshold (float, *optional*, 0.45): minimal overlap score for removing objects duplicates in NMS
agnostic_nms (bool, *optiona*, False): apply class agnostinc NMS approach or not
max_detections (int, *optional*, 300): maximum detections after NMS
pred_masks (np.ndarray, *optional*, None): model ooutput prediction masks, if not provided only boxes will be postprocessed
retina_mask (bool, *optional*, False): retina mask postprocessing instead of native decoding
Returns:
pred (List[Dict[str, np.ndarray]]): list of dictionary with det - detected boxes in format [x1, y1, x2, y2, score, label] and segment - segmentation polygons for each element in batch
"""
nms_kwargs = {"agnostic": agnosting_nms, "max_det":max_detections}
# if pred_masks is not None:
# nms_kwargs["nm"] = 32
preds = ops.non_max_suppression(
torch.from_numpy(pred_boxes),
min_conf_threshold,
nms_iou_threshold,
nc=5,
**nms_kwargs
)
results = []
proto = torch.from_numpy(pred_masks) if pred_masks is not None else None
for i, pred in enumerate(preds):
shape = orig_img[i].shape if isinstance(orig_img, list) else orig_img.shape
if not len(pred):
results.append({"det": [], "segment": []})
continue
if proto is None:
pred[:, :4] = ops.scale_boxes(input_hw, pred[:, :4], shape).round()
results.append({"det": pred})
continue
if retina_mask:
pred[:, :4] = ops.scale_boxes(input_hw, pred[:, :4], shape).round()
masks = ops.process_mask_native(proto[i], pred[:, 6:], pred[:, :4], shape[:2]) # HWC
segments = [ops.scale_segments(input_hw, x, shape, normalize=False) for x in ops.masks2segments(masks)]
else:
masks = ops.process_mask(proto[i], pred[:, 6:], pred[:, :4], input_hw, upsample=True)
pred[:, :4] = ops.scale_boxes(input_hw, pred[:, :4], shape).round()
segments = [ops.scale_segments(input_hw, x, shape, normalize=False) for x in ops.masks2segments(masks)]
results.append({"det": pred[:, :6].numpy(), "segment": segments})
return results
def detect(image:np.ndarray, model:Model):
"""
OpenVINO YOLOv8 model inference function. Preprocess image, runs model inference and postprocess results using NMS.
Parameters:
image (np.ndarray): input image.
model (Model): OpenVINO compiled model.
Returns:
detections (np.ndarray): detected boxes in format [x1, y1, x2, y2, score, label]
"""
num_outputs = len(model.outputs)
preprocessed_image = preprocess_image(image)
input_tensor = image_to_tensor(preprocessed_image)
start_time = time.time()
result = model(input_tensor)
end_time = time.time()
inference_time = (end_time - start_time) * 1000
# print(f"推理时长为：{inference_time}ms")
boxes = result[model.output(0)]
masks = None
if num_outputs > 1:
masks = result[model.output(1)]
input_hw = input_tensor.shape[2:]
detections = postprocess(pred_boxes=boxes, input_hw=input_hw, orig_img=image, pred_masks=masks)
return detections, inference_time
if __name__ == "__main__":
from ultralytics import YOLO
from tqdm import tqdm
import sys
import glob
seg_model_path = 'runs/segment/train6/weights/FP16_openvino_model/v8_seg_FP16.xml'
# seg_model_path = 'runs/segment/train6/weights/FP32_openvino_model/v8_seg_FP32.xml'
# seg_model_path = 'runs/segment/train6/weights/Int8_openvino_model/v8_seg_Int8.xml'
core = Core()
seg_ov_model = core.read_model(seg_model_path)
device = "CPU" # "GPU"
if device != "CPU":
seg_ov_model.reshape({0: [1, 3, 640, 640]})
seg_compiled_model = core.compile_model(seg_ov_model, device)
SEG_MODEL = YOLO(f'runs/segment/train6/weights/best.pt')
label_map = SEG_MODEL.model.names
image_folder = 'D:\\Ultralytics\\datasets\\coco128-seg\\images\\train2017\\*.jpg'
image_paths = glob.iglob(image_folder)
# Image_path = '2-12.13_RGB.jpg'
total_time = 0
index = 0
for image_path in image_paths:
sys.stdout.write('#')
sys.stdout.flush()
time.sleep(0.1)
# print(f"{(index/40) * 100}%")
input_image = np.array(Image.open(image_path))
detections, inference_time = detect(input_image, seg_compiled_model)
image_with_boxes= draw_results(detections[0], input_image, label_map)
b,g,r = cv2.split(image_with_boxes)
img_1 = cv2.merge([r,g,b])
index += 1
total_time += inference_time
cv2.namedWindow('test',cv2.WINDOW_NORMAL)
cv2.imshow('test', img_1)
cv2.waitKey(0)
# img = Image.fromarray(image_with_boxes)
#img.show()
print("\n",f"一共推理了{index}张图像，平均耗时：{total_time / index}ms")
问题以及解决方式
在评估代码中，有些同学可能会报ValidatorClass的bug，那是因为你的YOLO库中缺少了一些参数，因为我们下载的YOLO库不是最新版的原因，可以到github中的ultralytics项目的最新版本的开源代码中找到最新的YOLO库。这里我为了方便大家，我直接把代码贴在下面，将下面的代码替换到你项目中的YOLO库中即可：
class YOLO:
"""
YOLO (You Only Look Once) object detection model.
Args:
model (str, Path): Path to the model file to load or create.
type (str): Type/version of models to use. Defaults to "v8".
Attributes:
type (str): Type/version of models being used.
ModelClass (Any): Model class.
TrainerClass (Any): Trainer class.
ValidatorClass (Any): Validator class.
PredictorClass (Any): Predictor class.
predictor (Any): Predictor object.
model (Any): Model object.
trainer (Any): Trainer object.
task (str): Type of model task.
ckpt (Any): Checkpoint object if model loaded from *.pt file.
cfg (str): Model configuration if loaded from *.yaml file.
ckpt_path (str): Checkpoint file path.
overrides (dict): Overrides for trainer object.
metrics_data (Any): Data for metrics.
Methods:
__call__(): Alias for predict method.
_new(cfg, verbose=True): Initializes a new model and infers the task type from the model definitions.
_load(weights): Initializes a new model and infers the task type from the model head.
_check_is_pytorch_model(): Raises TypeError if model is not a PyTorch model.
reset(): Resets the model modules.
info(verbose=False): Logs model info.
fuse(): Fuse model for faster inference.
predict(source=None, stream=False, **kwargs): Perform prediction using the YOLO model.
Returns:
list(ultralytics.yolo.engine.results.Results): The prediction results.
"""
def __init__(self, model='yolov8n.pt', type='v8') -> None:
"""
Initializes the YOLO model.
Args:
model (str, Path): model to load or create
type (str): Type/version of models to use. Defaults to "v8".
"""
self._reset_callbacks()
self.type = type
self.ModelClass = None # model class
self.TrainerClass = None # trainer class
self.ValidatorClass = None # validator class
self.PredictorClass = None # predictor class
self.predictor = None # reuse predictor
self.model = None # model object
self.trainer = None # trainer object
self.task = None # task type
self.ckpt = None # if loaded from *.pt
self.cfg = None # if loaded from *.yaml
self.ckpt_path = None
self.overrides = {} # overrides for trainer object
self.metrics_data = None
# Load or create new YOLO model
suffix = Path(model).suffix
if not suffix and Path(model).stem in GITHUB_ASSET_STEMS:
model, suffix = Path(model).with_suffix('.pt'), '.pt' # add suffix, i.e. yolov8n -> yolov8n.pt
if suffix == '.yaml':
self._new(model)
else:
self._load(model)
def __call__(self, source=None, stream=False, **kwargs):
return self.predict(source, stream, **kwargs)
def _new(self, cfg: str, verbose=True):
"""
Initializes a new model and infers the task type from the model definitions.
Args:
cfg (str): model configuration file
verbose (bool): display model info on load
"""
self.cfg = check_yaml(cfg) # check YAML
cfg_dict = yaml_load(self.cfg, append_filename=True) # model dict
self.task = guess_model_task(cfg_dict)
self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = self._assign_ops_from_task()
self.model = self.ModelClass(cfg_dict, verbose=verbose and RANK == -1) # initialize
self.overrides['model'] = self.cfg
def _load(self, weights: str):
"""
Initializes a new model and infers the task type from the model head.
Args:
weights (str): model checkpoint to be loaded
"""
suffix = Path(weights).suffix
if suffix == '.pt':
self.model, self.ckpt = attempt_load_one_weight(weights)
self.task = self.model.args['task']
self.overrides = self.model.args
self._reset_ckpt_args(self.overrides)
self.ckpt_path = self.model.pt_path
else:
weights = check_file(weights)
self.model, self.ckpt = weights, None
self.task = guess_model_task(weights)
self.ckpt_path = weights
self.overrides['model'] = weights
self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = self._assign_ops_from_task()
def _check_is_pytorch_model(self):
"""
Raises TypeError is model is not a PyTorch model
"""
if not isinstance(self.model, nn.Module):
raise TypeError(f"model='{self.model}' must be a *.pt PyTorch model, but is a different type. "
f'PyTorch models can be used to train, val, predict and export, i.e. '
f"'yolo export model=yolov8n.pt', but exported formats like ONNX, TensorRT etc. only "
f"support 'predict' and 'val' modes, i.e. 'yolo predict model=yolov8n.onnx'.")
def reset(self):
"""
Resets the model modules.
"""
self._check_is_pytorch_model()
for m in self.model.modules():
if hasattr(m, 'reset_parameters'):
m.reset_parameters()
for p in self.model.parameters():
p.requires_grad = True
def info(self, verbose=False):
"""
Logs model info.
Args:
verbose (bool): Controls verbosity.
"""
self._check_is_pytorch_model()
self.model.info(verbose=verbose)
def fuse(self):
self._check_is_pytorch_model()
self.model.fuse()
def predict(self, source=None, stream=False, **kwargs):
"""
Perform prediction using the YOLO model.
Args:
source (str | int | PIL | np.ndarray): The source of the image to make predictions on.
Accepts all source types accepted by the YOLO model.
stream (bool): Whether to stream the predictions or not. Defaults to False.
**kwargs : Additional keyword arguments passed to the predictor.
Check the 'configuration' section in the documentation for all available options.
Returns:
(List[ultralytics.yolo.engine.results.Results]): The prediction results.
"""
overrides = self.overrides.copy()
overrides['conf'] = 0.25
overrides.update(kwargs)
overrides['mode'] = kwargs.get('mode', 'predict')
assert overrides['mode'] in ['track', 'predict']
overrides['save'] = kwargs.get('save', False) # not save files by default
if not self.predictor:
self.predictor = self.PredictorClass(overrides=overrides)
self.predictor.setup_model(model=self.model)
else: # only update args if predictor is already setup
self.predictor.args = get_cfg(self.predictor.args, overrides)
is_cli = sys.argv[0].endswith('yolo') or sys.argv[0].endswith('ultralytics')
return self.predictor.predict_cli(source=source) if is_cli else self.predictor(source=source, stream=stream)
@smart_inference_mode()
def track(self, source=None, stream=False, **kwargs):
from ultralytics.tracker import register_tracker
register_tracker(self)
# ByteTrack-based method needs low confidence predictions as input
conf = kwargs.get('conf') or 0.1
kwargs['conf'] = conf
kwargs['mode'] = 'track'
return self.predict(source=source, stream=stream, **kwargs)
@smart_inference_mode()
def val(self, data=None, **kwargs):
"""
Validate a model on a given dataset .
Args:
data (str): The dataset to validate on. Accepts all formats accepted by yolo
**kwargs : Any other args accepted by the validators. To see all args check 'configuration' section in docs
"""
overrides = self.overrides.copy()
overrides['rect'] = True # rect batches as default
overrides.update(kwargs)
overrides['mode'] = 'val'
args = get_cfg(cfg=DEFAULT_CFG, overrides=overrides)
args.data = data or args.data
args.task = self.task
if args.imgsz == DEFAULT_CFG.imgsz and not isinstance(self.model, (str, Path)):
args.imgsz = self.model.args['imgsz'] # use trained imgsz unless custom value is passed
args.imgsz = check_imgsz(args.imgsz, max_dim=1)
validator = self.ValidatorClass(args=args)
validator(model=self.model)
self.metrics_data = validator.metrics
return validator.metrics
@smart_inference_mode()
def benchmark(self, **kwargs):
"""
Benchmark a model on all export formats.
Args:
**kwargs : Any other args accepted by the validators. To see all args check 'configuration' section in docs
"""
from ultralytics.yolo.utils.benchmarks import run_benchmarks
overrides = self.model.args.copy()
overrides.update(kwargs)
overrides = {**DEFAULT_CFG_DICT, **overrides} # fill in missing overrides keys with defaults
return run_benchmarks(model=self, imgsz=overrides['imgsz'], half=overrides['half'], device=overrides['device'])
def export(self, **kwargs):
"""
Export model.
Args:
**kwargs : Any other args accepted by the predictors. To see all args check 'configuration' section in docs
"""
self._check_is_pytorch_model()
overrides = self.overrides.copy()
overrides.update(kwargs)
args = get_cfg(cfg=DEFAULT_CFG, overrides=overrides)
args.task = self.task
if args.imgsz == DEFAULT_CFG.imgsz:
args.imgsz = self.model.args['imgsz'] # use trained imgsz unless custom value is passed
if args.batch == DEFAULT_CFG.batch:
args.batch = 1 # default to 1 if not modified
exporter = Exporter(overrides=args)
return exporter(model=self.model)
def train(self, **kwargs):
"""
Trains the model on a given dataset.
Args:
**kwargs (Any): Any number of arguments representing the training configuration.
"""
self._check_is_pytorch_model()
overrides = self.overrides.copy()
overrides.update(kwargs)
if kwargs.get('cfg'):
LOGGER.info(f"cfg file passed. Overriding default params with {kwargs['cfg']}.")
overrides = yaml_load(check_yaml(kwargs['cfg']), append_filename=True)
overrides['task'] = self.task
overrides['mode'] = 'train'
if not overrides.get('data'):
raise AttributeError("Dataset required but missing, i.e. pass 'data=coco128.yaml'")
if overrides.get('resume'):
overrides['resume'] = self.ckpt_path
self.trainer = self.TrainerClass(overrides=overrides)
if not overrides.get('resume'): # manually set model only if not resuming
self.trainer.model = self.trainer.get_model(weights=self.model if self.ckpt else None, cfg=self.model.yaml)
self.model = self.trainer.model
self.trainer.train()
# update model and cfg after training
if RANK in {0, -1}:
self.model, _ = attempt_load_one_weight(str(self.trainer.best))
self.overrides = self.model.args
self.metrics_data = getattr(self.trainer.validator, 'metrics', None) # TODO: no metrics returned by DDP
def to(self, device):
"""
Sends the model to the given device.
Args:
device (str): device
"""
self._check_is_pytorch_model()
self.model.to(device)
def _assign_ops_from_task(self):
model_class, train_lit, val_lit, pred_lit = MODEL_MAP[self.task]
trainer_class = eval(train_lit.replace('TYPE', f'{self.type}'))
validator_class = eval(val_lit.replace('TYPE', f'{self.type}'))
predictor_class = eval(pred_lit.replace('TYPE', f'{self.type}'))
return model_class, trainer_class, validator_class, predictor_class
@property
def names(self):
"""
Returns class names of the loaded model.
"""
return self.model.names if hasattr(self.model, 'names') else None
@property
def device(self):
"""
Returns device if PyTorch model
"""
return next(self.model.parameters()).device if isinstance(self.model, nn.Module) else None
@property
def transforms(self):
"""
Returns transform of the loaded model.
"""
return self.model.transforms if hasattr(self.model, 'transforms') else None
@property
def metrics(self):
"""
Returns metrics if computed
"""
if not self.metrics_data:
LOGGER.info('No metrics data found! Run training or validation operation first.')
return self.metrics_data
@staticmethod
def add_callback(event: str, func):
"""
Add callback
"""
callbacks.default_callbacks[event].append(func)
@staticmethod
def _reset_ckpt_args(args):
for arg in 'augment', 'verbose', 'project', 'name', 'exist_ok', 'resume', 'batch', 'epochs', 'cache', \
'save_json', 'half', 'v5loader', 'device', 'cfg', 'save', 'rect', 'plots', 'opset', 'simplify':
args.pop(arg, None)
@staticmethod
def _reset_callbacks():
for event in callbacks.default_callbacks.keys():
callbacks.default_callbacks[event] = [callbacks.default_callbacks[event][0]]
3051
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
