69 lines
2.6 KiB
Python
69 lines
2.6 KiB
Python
import argparse
|
|
import os
|
|
import shutil
|
|
import csv
|
|
from ultralytics import YOLO
|
|
|
|
parser = argparse.ArgumentParser(description="YOLO26 岩性分类评估")
|
|
parser.add_argument("--model", default="runs/classify/train/weights/best.pt", help="模型权重路径")
|
|
parser.add_argument("--val-root", default="custom_dataset/val", help="验证集路径")
|
|
parser.add_argument("--save-root", default="岩性分类结果", help="分类结果保存目录")
|
|
parser.add_argument("--csv", default="分类预测结果.csv", help="预测结果 CSV 输出路径")
|
|
parser.add_argument("--device", default="cuda", help="推理设备")
|
|
args = parser.parse_args()
|
|
|
|
model = YOLO(args.model)
|
|
|
|
all_images = []
|
|
for cls_folder in os.listdir(args.val_root):
|
|
cls_path = os.path.join(args.val_root, cls_folder)
|
|
if os.path.isdir(cls_path):
|
|
for img_name in os.listdir(cls_path):
|
|
all_images.append(os.path.join(cls_path, img_name))
|
|
|
|
total_samples = 0
|
|
correct_samples = 0
|
|
|
|
with open(args.csv, "w", newline="", encoding="utf-8-sig") as f:
|
|
writer = csv.writer(f)
|
|
writer.writerow(["图片名", "真实标签", "预测类别", "置信度"])
|
|
|
|
for img_path in all_images:
|
|
res = model(img_path, verbose=False)[0]
|
|
pred_class = res.names[res.probs.top1]
|
|
conf = f"{res.probs.top1conf:.2%}"
|
|
|
|
true_class = os.path.basename(os.path.dirname(img_path))
|
|
img_name = os.path.basename(img_path)
|
|
|
|
writer.writerow([img_name, true_class, pred_class, conf])
|
|
|
|
target_dir = os.path.join(args.save_root, true_class)
|
|
os.makedirs(target_dir, exist_ok=True)
|
|
shutil.copy(img_path, os.path.join(target_dir, img_name))
|
|
|
|
total_samples += 1
|
|
|
|
is_correct = False
|
|
if true_class == "泥页岩" and ("泥" in pred_class or "页" in pred_class):
|
|
is_correct = True
|
|
elif true_class == "砂砾岩" and ("砂" in pred_class or "粉砂" in pred_class):
|
|
is_correct = True
|
|
elif true_class == "灰岩" and ("灰" in pred_class or "石灰岩" in pred_class):
|
|
is_correct = True
|
|
elif true_class == "云岩" and ("云" in pred_class or "白云" in pred_class):
|
|
is_correct = True
|
|
elif true_class == "特殊岩" and ("石膏" in pred_class or "煤" in pred_class or "膏盐" in pred_class):
|
|
is_correct = True
|
|
|
|
if is_correct:
|
|
correct_samples += 1
|
|
|
|
accuracy = (correct_samples / total_samples) * 100
|
|
print("\n" + "=" * 60)
|
|
print("分类完成！准确率统计结果：")
|
|
print(f"总测试样本数：{total_samples} 张")
|
|
print(f"预测正确样本数：{correct_samples} 张")
|
|
print(f"整体分类准确率：{accuracy:.2f}%")
|
|
print("=" * 60)
|