63 lines
2.1 KiB
Python
63 lines
2.1 KiB
Python
|
import os
|
|||
|
|
import cv2
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
|
|||
|
|
|
|||
|
|
class FullImageClassifier:
|
|||
|
|
"""封装 YOLO 分类模型，对整张图像进行二分类（无缺陷/有缺陷）"""
|
|||
|
|
|
|||
|
|
def __init__(self, model_path):
|
|||
|
|
if not os.path.exists(model_path):
|
|||
|
|
raise FileNotFoundError(f"模型文件不存在: {model_path}")
|
|||
|
|
|
|||
|
|
self.model = YOLO(model_path)
|
|||
|
|
print(f"[INFO] 成功加载 YOLO 分类模型: {model_path}")
|
|||
|
|
|
|||
|
|
def classify(self, img_np):
|
|||
|
|
"""
|
|||
|
|
对整张图像进行分类，返回类别 ID（0 或 1）
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
img_np (np.ndarray): BGR 格式的 OpenCV 图像 (H, W, C)
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
int: 分类结果，{0: "有缺陷", 1: "无缺陷"}
|
|||
|
|
"""
|
|||
|
|
# 直接 resize 整图到模型输入尺寸（YOLO 默认为 224x224 或 640x640，由训练决定）
|
|||
|
|
# Ultralytics YOLO 会自动处理 resize，但显式指定更可控
|
|||
|
|
resized_img = cv2.resize(img_np, (640, 640))
|
|||
|
|
|
|||
|
|
# 推理（verbose=False 关闭进度条）
|
|||
|
|
results = self.model.predict(resized_img, verbose=False)
|
|||
|
|
|
|||
|
|
cls = int(results[0].probs.top1) # 获取 top-1 类别索引
|
|||
|
|
return cls
|
|||
|
|
|
|||
|
|
|
|||
|
|
def cls_quexian():
|
|||
|
|
# ================== 配置 ==================
|
|||
|
|
model_path = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/cls/exp_xiantiao_cls/weights/best.pt"
|
|||
|
|
img_path = "1.png"
|
|||
|
|
|
|||
|
|
# ================== 1. 加载模型 ==================
|
|||
|
|
classifier = FullImageClassifier(model_path)
|
|||
|
|
|
|||
|
|
# ================== 2. 加载图像 ==================
|
|||
|
|
if not os.path.exists(img_path):
|
|||
|
|
raise FileNotFoundError(f"图片不存在: {img_path}")
|
|||
|
|
|
|||
|
|
img_np = cv2.imread(img_path)
|
|||
|
|
if img_np is None:
|
|||
|
|
raise ValueError("图像加载失败，可能路径错误或图像文件损坏。")
|
|||
|
|
|
|||
|
|
# ================== 3. 推理（整图）==================
|
|||
|
|
result = classifier.classify(img_np)
|
|||
|
|
|
|||
|
|
# ================== 4. 输出 ==================
|
|||
|
|
label_map = {0: "有缺陷", 1: "无缺陷"}
|
|||
|
|
print(f"\n===== 推理结果 =====")
|
|||
|
|
print(f"分类结果：{result} → {label_map.get(result, '未知')}\n")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
cls_quexian()
|