86 lines
2.9 KiB
Python
86 lines
2.9 KiB
Python
|
import os
|
|||
|
|
import cv2
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ROIClassifier:
|
|||
|
|
""" 封装 YOLO 分类模型 + ROI txt 推理 """
|
|||
|
|
|
|||
|
|
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 load_roi(self, roi_txt_path):
|
|||
|
|
"""读取 ROI txt，可自动支持逗号/空格分隔"""
|
|||
|
|
if not os.path.exists(roi_txt_path):
|
|||
|
|
raise FileNotFoundError(f"ROI 文件不存在: {roi_txt_path}")
|
|||
|
|
with open(roi_txt_path, "r") as f:
|
|||
|
|
text = f.read().strip()
|
|||
|
|
# 把逗号替换成空格，再 split
|
|||
|
|
for ch in [",", ";"]:
|
|||
|
|
text = text.replace(ch, " ")
|
|||
|
|
parts = text.split()
|
|||
|
|
if len(parts) != 4:
|
|||
|
|
raise ValueError(f"ROI txt 格式错误，应为4个数字，解析得到: {parts}\n文件: {roi_txt_path}")
|
|||
|
|
x1, y1, x2, y2 = map(int, parts)
|
|||
|
|
return x1, y1, x2, y2
|
|||
|
|
|
|||
|
|
def classify(self, img_np, roi_txt_path):
|
|||
|
|
"""对 ROI 区域做分类，返回 0/1"""
|
|||
|
|
h, w = img_np.shape[:2]
|
|||
|
|
x1, y1, x2, y2 = self.load_roi(roi_txt_path)
|
|||
|
|
|
|||
|
|
# -------- ROI 边界安全裁剪 --------
|
|||
|
|
x1 = max(0, min(x1, w - 1))
|
|||
|
|
x2 = max(0, min(x2, w - 1))
|
|||
|
|
y1 = max(0, min(y1, h - 1))
|
|||
|
|
y2 = max(0, min(y2, h - 1))
|
|||
|
|
|
|||
|
|
if x2 <= x1 or y2 <= y1:
|
|||
|
|
raise ValueError(f"ROI坐标无效: {x1, y1, x2, y2}")
|
|||
|
|
|
|||
|
|
# -------- 1. 裁剪 ROI --------
|
|||
|
|
roi_img = img_np[y1:y2, x1:x2]
|
|||
|
|
|
|||
|
|
# -------- 2. resize 到 640×640（强制送给模型）--------
|
|||
|
|
roi_img = cv2.resize(roi_img, (640, 640))
|
|||
|
|
|
|||
|
|
# -------- 3. YOLO 分类推理 --------
|
|||
|
|
results = self.model.predict(roi_img, verbose=False)
|
|||
|
|
|
|||
|
|
cls = int(results[0].probs.top1) # 0 或 1
|
|||
|
|
return cls
|
|||
|
|
|
|||
|
|
|
|||
|
|
def class_xiantiao():
|
|||
|
|
# ================== 配置 ==================
|
|||
|
|
model_path = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/cls/exp_xiantiao_cls/weights/best.pt"
|
|||
|
|
img_path = "1.png"
|
|||
|
|
roi_txt_path = "/home/hx/开发/ML_xiantiao/class_xiantiao/roi_1/1/1_rois1.txt"
|
|||
|
|
|
|||
|
|
# ================== 1. 加载模型 ==================
|
|||
|
|
classifier = ROIClassifier(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, roi_txt_path)
|
|||
|
|
|
|||
|
|
# ================== 4. 输出 ==================
|
|||
|
|
print(f"\n===== 推理结果 =====")
|
|||
|
|
print(f"ROI 文件：{roi_txt_path}")
|
|||
|
|
print(f"分类结果：{result} (0=异常 / 1=正常)\n")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
class_xiantiao()
|