169 lines
5.5 KiB
Python
169 lines
5.5 KiB
Python
|
import os
|
|||
|
|
from pathlib import Path
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
|
|||
|
|
# ---------------------------
|
|||
|
|
# 类别映射
|
|||
|
|
# ---------------------------
|
|||
|
|
CLASS_NAMES = {
|
|||
|
|
0: "未堆料",
|
|||
|
|
1: "小堆料",
|
|||
|
|
2: "大堆料",
|
|||
|
|
3: "未浇筑满",
|
|||
|
|
4: "浇筑满"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# ---------------------------
|
|||
|
|
# 加载 ROI 列表
|
|||
|
|
# ---------------------------
|
|||
|
|
def load_global_rois(txt_path):
|
|||
|
|
rois = []
|
|||
|
|
if not os.path.exists(txt_path):
|
|||
|
|
print(f"❌ ROI 文件不存在: {txt_path}")
|
|||
|
|
return rois
|
|||
|
|
with open(txt_path, 'r') as f:
|
|||
|
|
for line in f:
|
|||
|
|
s = line.strip()
|
|||
|
|
if s:
|
|||
|
|
try:
|
|||
|
|
x, y, w, h = map(int, s.split(','))
|
|||
|
|
rois.append((x, y, w, h))
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"无法解析 ROI 行 '{s}': {e}")
|
|||
|
|
return rois
|
|||
|
|
|
|||
|
|
# ---------------------------
|
|||
|
|
# 裁剪并 resize ROI
|
|||
|
|
# ---------------------------
|
|||
|
|
def crop_and_resize(img, rois, target_size=640):
|
|||
|
|
crops = []
|
|||
|
|
h_img, w_img = img.shape[:2]
|
|||
|
|
for i, (x, y, w, h) in enumerate(rois):
|
|||
|
|
if x < 0 or y < 0 or x + w > w_img or y + h > h_img:
|
|||
|
|
continue
|
|||
|
|
roi = img[y:y+h, x:x+w]
|
|||
|
|
roi_resized = cv2.resize(roi, (target_size, target_size), interpolation=cv2.INTER_AREA)
|
|||
|
|
crops.append((roi_resized, i))
|
|||
|
|
return crops
|
|||
|
|
|
|||
|
|
# ---------------------------
|
|||
|
|
# class1/class2 加权判断
|
|||
|
|
# ---------------------------
|
|||
|
|
def weighted_small_large(pred_probs, threshold=0.4, w1=0.3, w2=0.7):
|
|||
|
|
p1 = float(pred_probs[1])
|
|||
|
|
p2 = float(pred_probs[2])
|
|||
|
|
total = p1 + p2
|
|||
|
|
if total > 0:
|
|||
|
|
score = (w1 * p1 + w2 * p2) / total
|
|||
|
|
else:
|
|||
|
|
score = 0.0
|
|||
|
|
final_class = "大堆料" if score >= threshold else "小堆料"
|
|||
|
|
return final_class, score, p1, p2
|
|||
|
|
|
|||
|
|
# ---------------------------
|
|||
|
|
# 单张图片推理函数
|
|||
|
|
# ---------------------------
|
|||
|
|
def classify_image_weighted(image, model, threshold=0.4):
|
|||
|
|
results = model(image)
|
|||
|
|
pred_probs = results[0].probs.data.cpu().numpy().flatten()
|
|||
|
|
class_id = int(pred_probs.argmax())
|
|||
|
|
confidence = float(pred_probs[class_id])
|
|||
|
|
class_name = CLASS_NAMES.get(class_id, f"未知类别({class_id})")
|
|||
|
|
|
|||
|
|
# class1/class2 使用加权得分
|
|||
|
|
if class_id in [1, 2]:
|
|||
|
|
final_class, score, p1, p2 = weighted_small_large(pred_probs, threshold=threshold)
|
|||
|
|
else:
|
|||
|
|
final_class = class_name
|
|||
|
|
score = confidence
|
|||
|
|
p1 = float(pred_probs[1])
|
|||
|
|
p2 = float(pred_probs[2])
|
|||
|
|
|
|||
|
|
return final_class, score, p1, p2
|
|||
|
|
|
|||
|
|
# ---------------------------
|
|||
|
|
# 批量推理主函数
|
|||
|
|
# ---------------------------
|
|||
|
|
def batch_classify_images(model_path, input_folder, output_root, roi_file, target_size=640, threshold=0.5):
|
|||
|
|
# 加载模型
|
|||
|
|
model = YOLO(model_path)
|
|||
|
|
|
|||
|
|
# 确保输出根目录存在
|
|||
|
|
output_root = Path(output_root)
|
|||
|
|
output_root.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
# 为所有类别创建目录
|
|||
|
|
class_dirs = {}
|
|||
|
|
for name in CLASS_NAMES.values():
|
|||
|
|
d = output_root / name
|
|||
|
|
d.mkdir(exist_ok=True)
|
|||
|
|
class_dirs[name] = d
|
|||
|
|
|
|||
|
|
rois = load_global_rois(roi_file)
|
|||
|
|
if not rois:
|
|||
|
|
print("❌ 没有有效 ROI，退出")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 遍历图片
|
|||
|
|
for img_path in Path(input_folder).glob("*.*"):
|
|||
|
|
if img_path.suffix.lower() not in ['.jpg', '.jpeg', '.png', '.bmp', '.tif']:
|
|||
|
|
continue
|
|||
|
|
try:
|
|||
|
|
img = cv2.imread(str(img_path))
|
|||
|
|
if img is None:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
crops = crop_and_resize(img, rois, target_size)
|
|||
|
|
|
|||
|
|
for roi_resized, roi_idx in crops:
|
|||
|
|
final_class, score, p1, p2 = classify_image_weighted(roi_resized, model, threshold=threshold)
|
|||
|
|
|
|||
|
|
# 文件名中保存 ROI、类别、加权分数、class1/class2 置信度
|
|||
|
|
suffix = f"_roi{roi_idx}_{final_class}_score{score:.2f}_p1{p1:.2f}_p2{p2:.2f}"
|
|||
|
|
dst_path = class_dirs[final_class] / f"{img_path.stem}{suffix}{img_path.suffix}"
|
|||
|
|
cv2.imwrite(dst_path, roi_resized)
|
|||
|
|
print(f"{img_path.name}{suffix} -> {final_class} (score={score:.2f}, p1={p1:.2f}, p2={p2:.2f})")
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"处理失败 {img_path.name}: {e}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------
|
|||
|
|
# 单张图片使用示例（保留 ROI，不保存文件）
|
|||
|
|
# ---------------------------
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
model_path = r"best.pt"
|
|||
|
|
image_path = r"./test_image/2.jpg" # 单张图片路径
|
|||
|
|
roi_file = r"./roi_coordinates/1_rois.txt"
|
|||
|
|
target_size = 640
|
|||
|
|
threshold = 0.4 #加权得分阈值可以根据大小堆料分类结果进行调整
|
|||
|
|
|
|||
|
|
# 加载模型
|
|||
|
|
model = YOLO(model_path)
|
|||
|
|
|
|||
|
|
# 读取 ROI
|
|||
|
|
rois = load_global_rois(roi_file)
|
|||
|
|
if not rois:
|
|||
|
|
print("❌ 没有有效 ROI，退出")
|
|||
|
|
exit(1)
|
|||
|
|
|
|||
|
|
# 读取图片
|
|||
|
|
img = cv2.imread(image_path)
|
|||
|
|
if img is None:
|
|||
|
|
print(f"❌ 无法读取图片: {image_path}")
|
|||
|
|
exit(1)
|
|||
|
|
|
|||
|
|
# 注意：必须裁剪 ROI 并推理，因为训练的时候输入的图像是经过resize的
|
|||
|
|
crops = crop_and_resize(img, rois, target_size)
|
|||
|
|
for roi_resized, roi_idx in crops:
|
|||
|
|
#final_class, score, p1, p2 = classify_image_weighted(roi_resized, model, threshold=threshold)
|
|||
|
|
final_class,_,_,_ = classify_image_weighted(roi_resized, model, threshold=threshold)
|
|||
|
|
# 只输出信息，不保存文件
|
|||
|
|
#print(f"ROI {roi_idx} -> 类别: {final_class}, 加权分数: {score:.2f}, "
|
|||
|
|
#f"class1 置信度: {p1:.2f}, class2 置信度: {p2:.2f}")
|
|||
|
|
print(f"类别: {final_class}")
|
|||
|
|
|
|||
|
|
|