# 导入库
import cv2
import numpy as np
from ultralytics import YOLO
from tensorflow.keras.models import load_model
import matplotlib.pyplot as plt
# 加载YOLO检测模型
yolo_model = YOLO("yolov8s.pt") # 替换为你的YOLO模型路径
# 加载分类模型,你可以运行notebook并保存模型并使用它(查看步骤2)
classification_model = load_model('dog_classification_model.h5')
# 分类标签
species_list = ['afghan_hound', 'african_hunting_dog', 'airedale', 'basenji', 'basset', 'beagle',
'bedlington_terrier', 'bernese_mountain_dog', 'black-and-tan_coonhound',
'blenheim_spaniel', 'bloodhound', 'bluetick', 'border_collie', 'border_terrier',
'borzoi', 'boston_bull', 'bouvier_des_flandres', 'brabancon_griffon', 'bull_mastiff',
'cairn', 'cardigan', 'chesapeake_bay_retriever', 'chow', 'clumber', 'cocker_spaniel',
'collie', 'curly-coated_retriever', 'dhole', 'dingo', 'doberman', 'english_foxhound',
'english_setter', 'entlebucher', 'flat-coated_retriever', 'german_shepherd',
'german_short-haired_pointer', 'golden_retriever', 'gordon_setter', 'great_dane',
'great_pyrenees', 'groenendael', 'ibizan_hound', 'irish_setter', 'irish_terrier',
'irish_water_spaniel', 'irish_wolfhound', 'japanese_spaniel', 'keeshond',
'kerry_blue_terrier', 'komondor', 'kuvasz', 'labrador_retriever', 'leonberg',
'lhasa', 'malamute', 'malinois', 'maltese_dog', 'mexican_hairless', 'miniature_pinscher',
'miniature_schnauzer', 'newfoundland', 'norfolk_terrier', 'norwegian_elkhound',
'norwich_terrier', 'old_english_sheepdog', 'otterhound', 'papillon', 'pekinese',
'pembroke', 'pomeranian', 'pug', 'redbone', 'rhodesian_ridgeback', 'rottweiler',
'saint_bernard', 'saluki', 'samoyed', 'schipperke', 'scotch_terrier',
'scottish_deerhound', 'sealyham_terrier', 'shetland_sheepdog', 'standard_poodle',
'standard_schnauzer', 'sussex_spaniel', 'tibetan_mastiff', 'tibetan_terrier',
'toy_terrier', 'vizsla', 'weimaraner', 'whippet', 'wire-haired_fox_terrier',
'yorkshire_terrier']
# 执行推理
def classify_region(image, model, target_size=(180, 180)): # 尺寸必须与分类模型的输入匹配
input_image = preprocess_image(image, target_size)
predictions = model.predict(input_image)
predicted_index = np.argmax(predictions[0])
predicted_label = species_list[predicted_index]
return predicted_label
# 加载图像
image_path = r"test-images/dog12.jpg" # 图像路径
image = cv2.imread(image_path)
# YOLO推理 --> 目标检测模型
results = yolo_model(image)
detections = results[0].boxes # 获取检测结果
# 检查YOLO的标签是否为"dog"并处理边界框
for detection in detections:
x1, y1, x2, y2 = map(int, detection.xyxy[0].tolist()) # 获取边界框坐标
conf = float(detection.conf[0]) # 获取置信度
cls_label = yolo_model.names[int(detection.cls[0])] # 直接从YOLO获取标签名称
# 检查标签是否为"dog"
if cls_label == "dog":
"""
提取用于分类的感兴趣区域(ROI)。
记住,图像分类模型只会对检测到的对象进行处理,
而不是对整个图像进行处理。
"""
roi = image[y1:y2, x1:x2]
# 如果ROI足够大,则对其进行分类
if roi.shape[0] > 0 and roi.shape[1] > 0:
# 图像分类模型
label = classify_region(roi, classification_model)
bbox_height = y2 - y1
font_scale = bbox_height / 200.0 # 比例因子,可根据需要调整
font_thickness = max(1, int(bbox_height / 100)) # 确保厚度至少为1
# 绘制边界框和标签
cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 4)
cv2.putText(image, label, (x1+100, y1-20), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 255), font_thickness)
print(f"检测到的狗品种: {label}")
cv2.imwrite("dog2-result.jpg", image)
# 显示结果图像
plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
plt.axis("off")
plt.show()
- 1.
- 2.
- 3.
- 4.
- 5.
- 6.
- 7.
- 8.
- 9.
- 10.
- 11.
- 12.
- 13.
- 14.
- 15.
- 16.
- 17.
- 18.
- 19.
- 20.
- 21.
- 22.
- 23.
- 24.
- 25.
- 26.
- 27.
- 28.
- 29.
- 30.
- 31.
- 32.
- 33.
- 34.
- 35.
- 36.
- 37.
- 38.
- 39.
- 40.
- 41.
- 42.
- 43.
- 44.
- 45.
- 46.
- 47.
- 48.
- 49.
- 50.
- 51.
- 52.
- 53.
- 54.
- 55.
- 56.
- 57.
- 58.
- 59.
- 60.
- 61.
- 62.
- 63.
- 64.
- 65.
- 66.
- 67.
- 68.
- 69.
- 70.
- 71.
- 72.
- 73.
- 74.
- 75.
- 76.
- 77.
- 78.
- 79.
- 80.
- 81.
- 82.
- 83.
- 84.
- 85.
- 86.
- 87.