模型量化与加载模版

1.模型量化
from ultralytics import YOLO

if name == 'main':

model_path = r"runs/detect/train/weights/best.pt"
data_path = r"D:\01-AI\20260914\dataset\images\images\YOLODataset\dataset.yaml"
model = YOLO(model_path)
model.export(format="onnx", int8=True, data=data_path) # 量化为int8

model1 = YOLO("runs/detect/train/weights/best_int8.onnx")

model1.val(data=data_path)

2.1模型加载（函数加载）
from ultralytics import YOLO
import cv2

img_path = "01.jpg"
img = cv2.imread(img_path)
if img is None:
raise FileNotFoundError(f"无法读取图片: {img_path}")

加载 ONNX 模型（而非 .pt）

model = YOLO("runs/detect/train/weights/best_int8.onnx")

正常调用 predict

results = model.predict(source=img_path,
conf=0.25, iou=0.45,
device='cpu' # 设备使用cpu推理即可
)

所有的类别名称

names = model.names

结果处理和 .pt 模型完全一致

boxes = results[0].boxes
print(boxes.xyxy, boxes.conf, boxes.cls)
for x1, y1, x2, y2, score, cls in boxes.data:
p1, p2 = (int(x1), int(y1)), (int(x2), int(y2))
cv2.rectangle(img, p1, p2, (0, 255, 0), 2)
label = f"{names[int(cls)]} {score:.2f}"
cv2.putText(img, label, (p1[0], max(p1[1] - 5, 0)),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
cv2.imwrite("test.jpg", img)
print(f"[INFO] 检测到 {len(boxes)} 个目标，已保存至 test.png")

2.1模型加载（自定义加载）
import cv2
import numpy as np
import onnxruntime as ort

ONNX_PATH = "runs/detect/train/weights/best_int8.onnx"
IMAGE_PATH = "01.jpg"
OUTPUT_PATH = "test.jpg"
CONF, IOU, SIZE = 0.25, 0.45, 640

---------- 1. 加载模型 ----------

session = ort.InferenceSession(ONNX_PATH, providers=["CPUExecutionProvider"])
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

---------- 2. 预处理 ----------

def preprocess(img):
h0, w0 = img.shape[:2]
r = min(SIZE / h0, SIZE / w0)
new_w, new_h = int(w0 * r), int(h0 * r)

resized = cv2.resize(img, (new_w, new_h))
dw, dh = (SIZE - new_w) // 2, (SIZE - new_h) // 2
padded = cv2.copyMakeBorder(resized, dh, SIZE - new_h - dh,
dw, SIZE - new_w - dw,
cv2.BORDER_CONSTANT, value=(114, 114, 114))

blob = padded[:, :, ::-1].transpose(2, 0, 1)[None] / 255.0
return blob.astype(np.float32), r, dw, dh

---------- 3. 后处理 ----------

def xywh2xyxy(x):
y = np.empty_like(x)
y[:, 0] = x[:, 0] - x[:, 2] / 2
y[:, 1] = x[:, 1] - x[:, 3] / 2
y[:, 2] = x[:, 0] + x[:, 2] / 2
y[:, 3] = x[:, 1] + x[:, 3] / 2
return y

def nms(boxes, scores):
order = scores.argsort()[::-1]
keep = []
while order.size > 0:
i = order[0]
keep.append(i)

xx1 = np.maximum(boxes[i, 0], boxes[order[1:], 0])
yy1 = np.maximum(boxes[i, 1], boxes[order[1:], 1])
xx2 = np.minimum(boxes[i, 2], boxes[order[1:], 2])
yy2 = np.minimum(boxes[i, 3], boxes[order[1:], 3])

w = np.maximum(0, xx2 - xx1)
h = np.maximum(0, yy2 - yy1)
inter = w * h
area_i = (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
area_o = (boxes[order[1:], 2] - boxes[order[1:], 0]) *
(boxes[order[1:], 3] - boxes[order[1:], 1])
iou = inter / (area_i + area_o - inter + 1e-7)

order = order[np.where(iou <= IOU)[0] + 1]
return keep

def postprocess(pred, r, dw, dh, h0, w0):
pred = pred[0].T # [8400, 4+nc]
boxes, scores_all = pred[:, :4], pred[:, 4:]

cls = np.argmax(scores_all, axis=1)
scores = scores_all[np.arange(len(cls)), cls]

mask = scores > CONF
boxes, scores, cls = boxes[mask], scores[mask], cls[mask]
if len(boxes) == 0:
return np.zeros((0, 6))

还原到原图坐标

boxes = xywh2xyxy(boxes)
boxes[:, [0, 2]] = (boxes[:, [0, 2]] - dw) / r
boxes[:, [1, 3]] = (boxes[:, [1, 3]] - dh) / r
boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w0)
boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h0)

keep = nms(boxes, scores)
return np.concatenate([boxes[keep], scores[keep, None], cls[keep, None]], axis=1)

if name == 'main':

---------- 4. 主流程 ----------

img = cv2.imread(IMAGE_PATH)
h0, w0 = img.shape[:2]

blob, r, dw, dh = preprocess(img)
pred = session.run([output_name], {input_name: blob})[0]
dets = postprocess(pred, r, dw, dh, h0, w0)

for x1, y1, x2, y2, score, cls in dets:
cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
cv2.putText(img, f"{int(cls)} {score:.2f}", (int(x1), int(y1) - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

cv2.imwrite(OUTPUT_PATH, img)
print(f"检测到 {len(dets)} 个目标，已保存至 {OUTPUT_PATH}")