69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
|
import os
|
||
|
|
from pathlib import Path
|
||
|
|
import cv2
|
||
|
|
from ultralytics import YOLO
|
||
|
|
|
||
|
|
# ---------------------------
|
||
|
|
# 配置路径（请按需修改）
|
||
|
|
# ---------------------------
|
||
|
|
MODEL_PATH = "gaiban.pt" # 你的二分类模型
|
||
|
|
INPUT_FOLDER = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/12.2" # 输入图像文件夹
|
||
|
|
OUTPUT_ROOT = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/12.2.2" # 输出根目录（会生成 合格/不合格 子文件夹）
|
||
|
|
|
||
|
|
# 类别映射（必须与训练时的 data.yaml 顺序一致）
|
||
|
|
CLASS_NAMES = {0: "不合格", 1: "合格"}
|
||
|
|
|
||
|
|
# ---------------------------
|
||
|
|
# 批量推理函数
|
||
|
|
# ---------------------------
|
||
|
|
def batch_classify(model_path, input_folder, output_root):
|
||
|
|
# 加载模型
|
||
|
|
model = YOLO(model_path)
|
||
|
|
print(f"✅ 模型加载成功: {model_path}")
|
||
|
|
|
||
|
|
# 创建输出目录
|
||
|
|
output_root = Path(output_root)
|
||
|
|
for cls_name in CLASS_NAMES.values():
|
||
|
|
(output_root / cls_name).mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
# 支持的图像格式
|
||
|
|
IMG_EXTS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}
|
||
|
|
|
||
|
|
input_dir = Path(input_folder)
|
||
|
|
processed = 0
|
||
|
|
|
||
|
|
for img_path in input_dir.iterdir():
|
||
|
|
if img_path.suffix.lower() not in IMG_EXTS:
|
||
|
|
continue
|
||
|
|
|
||
|
|
# 读取图像
|
||
|
|
img = cv2.imread(str(img_path))
|
||
|
|
if img is None:
|
||
|
|
print(f"❌ 无法读取: {img_path}")
|
||
|
|
continue
|
||
|
|
|
||
|
|
# 推理（整图）
|
||
|
|
results = model(img)
|
||
|
|
probs = results[0].probs.data.cpu().numpy()
|
||
|
|
pred_class_id = int(probs.argmax())
|
||
|
|
pred_label = CLASS_NAMES[pred_class_id]
|
||
|
|
confidence = float(probs[pred_class_id])
|
||
|
|
|
||
|
|
# 保存原图到对应文件夹
|
||
|
|
dst = output_root / pred_label / img_path.name
|
||
|
|
cv2.imwrite(str(dst), img)
|
||
|
|
|
||
|
|
print(f"✅ {img_path.name} → {pred_label} ({confidence:.2f})")
|
||
|
|
processed += 1
|
||
|
|
|
||
|
|
print(f"\n🎉 共处理 {processed} 张图像，结果已保存至: {output_root}")
|
||
|
|
|
||
|
|
# ---------------------------
|
||
|
|
# 运行入口
|
||
|
|
# ---------------------------
|
||
|
|
if __name__ == "__main__":
|
||
|
|
batch_classify(
|
||
|
|
model_path=MODEL_PATH,
|
||
|
|
input_folder=INPUT_FOLDER,
|
||
|
|
output_root=OUTPUT_ROOT
|
||
|
|
)
|