133 lines
4.0 KiB
Python
133 lines
4.0 KiB
Python
import os
|
||
from pathlib import Path
|
||
import cv2
|
||
import numpy as np
|
||
from ultralytics import YOLO
|
||
import shutil
|
||
|
||
# ---------------------------
|
||
# 全局配置
|
||
# ---------------------------
|
||
|
||
# ✅ 定义唯一的 ROI (x, y, w, h) —— 请根据你的实际坐标修改！
|
||
SINGLE_ROI = (859,810,696,328) # 格式: x, y, w, h
|
||
|
||
CLASS_NAMES = {
|
||
0: "未堆料",
|
||
1: "小堆料",
|
||
2: "大堆料",
|
||
3: "未浇筑满",
|
||
4: "浇筑满"
|
||
}
|
||
|
||
TARGET_SIZE = 640
|
||
THRESHOLD = 0.4 # 加权得分阈值
|
||
|
||
|
||
# ---------------------------
|
||
# 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
|
||
|
||
|
||
# ---------------------------
|
||
# 批量推理主函数（单 ROI + 批量处理）
|
||
# ---------------------------
|
||
def batch_classify_images(model_path, input_folder, output_root, target_size=640, threshold=0.4):
|
||
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
|
||
|
||
x, y, w, h = SINGLE_ROI
|
||
|
||
# 存储所有裁剪后的图像和对应的原始路径
|
||
crops = []
|
||
img_paths = []
|
||
|
||
input_folder = Path(input_folder)
|
||
supported_exts = {'.jpg', '.jpeg', '.png', '.bmp', '.tif'}
|
||
|
||
print("🔍 正在裁剪所有图片的 ROI 区域...")
|
||
for img_path in input_folder.glob("*.*"):
|
||
if img_path.suffix.lower() not in supported_exts:
|
||
continue
|
||
img = cv2.imread(str(img_path))
|
||
if img is None:
|
||
continue
|
||
|
||
# 裁剪并 resize 到模型输入大小
|
||
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)
|
||
img_paths.append(img_path)
|
||
|
||
if not crops:
|
||
print("❌ 没有有效图片可供处理")
|
||
return
|
||
|
||
total = len(crops)
|
||
print(f"✅ 共准备 {total} 张图片，开始批量推理...")
|
||
|
||
# 🔥 批量推理（一次 forward pass）
|
||
results = model(
|
||
source=crops,
|
||
verbose=False,
|
||
imgsz=target_size,
|
||
half=False, # 如果使用 GPU 可开启半精度: half=True
|
||
device=0 # 使用 GPU 0，如用 CPU 改为: device=None
|
||
)
|
||
|
||
# 后处理结果
|
||
print("📦 正在处理结果并保存...")
|
||
for i, result in enumerate(results):
|
||
pred_probs = result.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})")
|
||
|
||
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])
|
||
|
||
# 构造目标路径
|
||
original_path = img_paths[i]
|
||
suffix = f"_roi0_{final_class}_score{score:.2f}_p1{p1:.2f}_p2{p2:.2f}"
|
||
dst_path = class_dirs[final_class] / f"{original_path.stem}{suffix}{original_path.suffix}"
|
||
|
||
# 复制原图到对应类别文件夹
|
||
shutil.copy2(str(original_path), str(dst_path))
|
||
print(f"{original_path.name} -> {final_class} (score={score:.2f}, p1={p1:.2f}, p2={p2:.2f})")
|
||
|
||
print(f"\n🎉 分类完成！共处理 {total} 张图片")
|
||
|
||
|
||
# ---------------------------
|
||
# 使用示例
|
||
# ---------------------------
|
||
if __name__ == "__main__":
|
||
MODEL_PATH = r"cls5.pt"
|
||
INPUT_FOLDER = r"./test_image"
|
||
OUTPUT_ROOT = r"./classified_images"
|
||
|
||
batch_classify_images(MODEL_PATH, INPUT_FOLDER, OUTPUT_ROOT, TARGET_SIZE, THRESHOLD) |