89 lines
2.6 KiB
Python
89 lines
2.6 KiB
Python
import cv2
|
||
from ultralytics import YOLO
|
||
|
||
# ---------------------------
|
||
# 类别映射（必须与训练时 data.yaml 一致）
|
||
# ---------------------------
|
||
CLASS_NAMES = {
|
||
0: "未堆料",
|
||
1: "小堆料",
|
||
2: "大堆料",
|
||
3: "未浇筑满",
|
||
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
|
||
|
||
|
||
# ---------------------------
|
||
# 单张图片推理主函数
|
||
# ---------------------------
|
||
def classify_single_image(model_path, image_path, threshold=0.5):
|
||
# 加载模型
|
||
print("🚀 加载模型...")
|
||
model = YOLO(model_path)
|
||
|
||
# 读取图像
|
||
img = cv2.imread(image_path)
|
||
if img is None:
|
||
raise FileNotFoundError(f"❌ 无法读取图像: {image_path}")
|
||
|
||
print(f"📷 推理图像: {image_path}")
|
||
|
||
# 整图分类（不裁剪）
|
||
results = model(img)
|
||
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})")
|
||
|
||
# 对 小堆料/大堆料 使用加权逻辑
|
||
if class_id in [1, 2]:
|
||
final_class, score, p1, p2 = weighted_small_large(pred_probs, threshold=threshold)
|
||
print("\n🔍 检测到堆料区域，使用加权判断:")
|
||
print(f" 小堆料概率: {p1:.4f}")
|
||
print(f" 大堆料概率: {p2:.4f}")
|
||
print(f" 加权得分: {score:.4f} (阈值={threshold})")
|
||
else:
|
||
final_class = class_name
|
||
score = confidence
|
||
|
||
# 输出最终结果
|
||
print("\n" + "="*40)
|
||
print(f"最终分类结果: {final_class}")
|
||
print(f"置信度/得分: {score:.4f}")
|
||
print("="*40)
|
||
|
||
return final_class, score
|
||
|
||
|
||
# ---------------------------
|
||
# 运行入口（请修改路径）
|
||
# ---------------------------
|
||
if __name__ == "__main__":
|
||
MODEL_PATH = "61best.pt"
|
||
IMAGE_PATH = "class4.png" # 👈 改成你的单张图片路径
|
||
|
||
# 可选：调整加权阈值（默认 0.4）
|
||
THRESHOLD = 0.4
|
||
|
||
try:
|
||
result_class, result_score = classify_single_image(
|
||
model_path=MODEL_PATH,
|
||
image_path=IMAGE_PATH,
|
||
threshold=THRESHOLD
|
||
)
|
||
except Exception as e:
|
||
print(f"程序出错: {e}") |