131 lines
4.3 KiB
Python
131 lines
4.3 KiB
Python
import os
|
||
from pathlib import Path
|
||
import cv2
|
||
import shutil
|
||
from ultralytics import YOLO
|
||
|
||
# ---------------------------
|
||
# 三分类类别定义（必须与模型训练时的顺序一致！）
|
||
# ---------------------------
|
||
CLASS_NAMES = {
|
||
0: "模具车 1",
|
||
1: "模具车 2",
|
||
2: "有遮挡"
|
||
}
|
||
|
||
|
||
# ---------------------------
|
||
# 单张图片推理函数（直接输入原图）
|
||
# ---------------------------
|
||
def classify_full_image(image_numpy, model):
|
||
"""
|
||
直接将整张原图送入模型推理
|
||
输入：numpy 数组 (BGR)
|
||
输出：(类别名称，置信度)
|
||
"""
|
||
# YOLO classification 模型会自动将输入图像 resize 到训练时的大小 (如 224x224 或 640x640)
|
||
results = model(image_numpy)
|
||
|
||
# 获取概率分布
|
||
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})")
|
||
return class_name, confidence
|
||
|
||
|
||
# ---------------------------
|
||
# 批量推理主函数 (直接推理原图并移动)
|
||
# ---------------------------
|
||
def batch_classify_full_images(model_path, input_folder, output_root):
|
||
print(f"🚀 开始加载模型：{model_path}")
|
||
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
|
||
print(f"✅ 准备输出目录：{d}")
|
||
|
||
input_folder = Path(input_folder)
|
||
image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff'}
|
||
processed_count = 0
|
||
error_count = 0
|
||
|
||
# 获取所有图片文件
|
||
image_files = [f for f in input_folder.iterdir() if f.suffix.lower() in image_extensions]
|
||
print(f"\n📂 发现 {len(image_files)} 张图片，开始全图推理...\n")
|
||
|
||
for img_path in image_files:
|
||
try:
|
||
print(f"📄 处理：{img_path.name}")
|
||
|
||
# 1. 读取原图
|
||
img = cv2.imread(str(img_path))
|
||
if img is None:
|
||
print(f"❌ 无法读取图像 (可能是损坏或格式不支持): {img_path.name}")
|
||
error_count += 1
|
||
continue
|
||
|
||
# 2. 【核心变化】直接对整张图进行推理，不再裁剪 ROI
|
||
final_class, conf = classify_full_image(img, model)
|
||
print(f" 🔍 全图识别结果：{final_class} (conf={conf:.2f})")
|
||
|
||
# 3. 确定目标目录
|
||
dst_dir = class_dirs[final_class]
|
||
dst_path = dst_dir / img_path.name
|
||
|
||
# 处理文件名冲突 (如果目标文件夹已有同名文件)
|
||
if dst_path.exists():
|
||
stem = img_path.stem
|
||
suffix = img_path.suffix
|
||
counter = 1
|
||
while True:
|
||
new_name = f"{stem}_{counter}{suffix}"
|
||
dst_path = dst_dir / new_name
|
||
if not dst_path.exists():
|
||
break
|
||
counter += 1
|
||
print(f" ⚠️ 目标文件已存在，重命名为：{dst_path.name}")
|
||
|
||
# 4. 移动【原图】到对应分类文件夹
|
||
shutil.move(str(img_path), str(dst_path))
|
||
print(f" ✅ 成功移动原图 -> [{final_class}]")
|
||
|
||
processed_count += 1
|
||
|
||
except Exception as e:
|
||
print(f"❌ 处理失败 {img_path.name}: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
error_count += 1
|
||
|
||
print(f"\n🎉 批量处理完成！")
|
||
print(f" 成功移动：{processed_count} 张")
|
||
print(f" 失败/跳过：{error_count} 张")
|
||
|
||
|
||
# ---------------------------
|
||
# 主程序入口
|
||
# ---------------------------
|
||
if __name__ == "__main__":
|
||
# 配置路径
|
||
model_path = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/cls_resize_muju/exp_cls2/weights/best.pt"
|
||
input_folder = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/61/浇筑满"
|
||
output_root = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/61"
|
||
|
||
# 注意：不再需要 roi_file 和 target_size 参数
|
||
|
||
batch_classify_full_images(
|
||
model_path=model_path,
|
||
input_folder=input_folder,
|
||
output_root=output_root
|
||
) |