import argparse
from ultralytics import YOLO
import matplotlib.pyplot as plt

parser = argparse.ArgumentParser(description="YOLO26 单张图片岩性分类预测")
parser.add_argument("--image", required=True, help="待预测的图片路径")
parser.add_argument("--model", default="runs/classify/train/weights/best.pt", help="模型权重路径")
parser.add_argument("--output", default="rock_result.png", help="结果图片保存路径")
parser.add_argument("--device", default="cuda", help="推理设备")
args = parser.parse_args()

model = YOLO(args.model)

results = model(args.image, device=args.device)

print("=" * 50)
print("岩石分类预测结果")
print(f"图片：{args.image}")
print(f"类别：{results[0].names[results[0].probs.top1]}")
print(f"置信度：{results[0].probs.top1conf:.2%}")
print("=" * 50)

plt.imshow(results[0].plot())
plt.axis("off")
plt.savefig(args.output, dpi=300, bbox_inches="tight")
plt.show()