解決方案:結(jié)合 YOLO 和 TensorFlow 做目標檢測和圖像分類
目標檢測和圖像分類是兩個不同的任務(wù),每個任務(wù)都有其特定的用途。在本文中,我將向你解釋什么是目標檢測和圖像分類,如何訓(xùn)練模型,最后,我將同時使用目標檢測和圖像分類模型來分類和檢測狗的品種。
目標檢測 + 圖像分類
目標檢測
目標檢測是一項基本的計算機視覺任務(wù),用于檢測和定位物體。簡而言之,目標檢測模型接受圖像作為輸入,并輸出坐標和標簽。
目標檢測示例
如果你需要知道物體的坐標,那么你需要使用目標檢測模型,因為圖像分類模型不會輸出坐標,它們只返回標簽。
圖像分類
圖像分類只輸出標簽。你將圖像作為輸入提供給模型,它將標簽作為輸出返回。它更適合分類相同類型的物體。例如,如果你想分類海洋動物,你需要訓(xùn)練一個圖像分類模型。
海洋動物分類
為什么不只使用目標檢測模型?
你可能已經(jīng)注意到,目標檢測模型同時提供坐標和標簽作為輸出,你可能會對自己說,為什么不直接使用目標檢測模型來處理所有事情呢?畢竟,它們理論上同時給出坐標和標簽,所以不需要分類模型。你可能一開始會這樣想,但有一些不同的因素你可能沒有意識到:
目標檢測模型非常適合識別和定位場景中的各種物體。但當(dāng)涉及到區(qū)分幾乎相同的物體時,圖像分類模型通常表現(xiàn)更好(一般來說,并非總是如此)。
你并不總能找到合適的數(shù)據(jù)集,創(chuàng)建數(shù)據(jù)集可能既耗時又無聊。如果你決定創(chuàng)建自己的數(shù)據(jù)集,創(chuàng)建圖像分類數(shù)據(jù)集比創(chuàng)建目標檢測數(shù)據(jù)集要容易得多。
目標檢測 + 圖像分類
工作流程
首先,我們將使用yolov8目標檢測模型檢測物體,然后從檢測到的物體中,我們將嘗試使用圖像分類模型對它們進行分類。請注意,圖像分類模型將僅對檢測到的物體進行操作,而不是對整個圖像進行操作。
1. 用于檢測狗的目標檢測模型
首先,可以參考文章《定制YOLOv8模型 - 檢測棋盤棋子》訓(xùn)練一個YOLO目標檢測模型。現(xiàn)在,我將使用一個預(yù)訓(xùn)練的YOLOv8模型,因為它包含狗類,我將直接使用這個預(yù)訓(xùn)練模型。我將使用YOLO模型進行檢測,如果它檢測到狗,我將使用圖像分類模型繼續(xù)處理。一般來說,最好為特定任務(wù)使用特定數(shù)據(jù)集訓(xùn)練模型。
2. 用于狗品種的圖像分類模型
我將使用TensorFlow來訓(xùn)練圖像分類模型。訓(xùn)練模型可能需要時間,具體取決于數(shù)據(jù)集和參數(shù)。有關(guān)狗品種分類模型的詳細代碼可以查看鏈接:https://www.kaggle.com/code/merfarukgnaydn/dog-species-classification
使用TensorFlow Keras訓(xùn)練圖像分類模型
3. 結(jié)合目標檢測 + 圖像分類模型
正如我之前向你解釋的那樣,這個過程非常簡單。首先,目標檢測模型執(zhí)行,然后是圖像分類模型。下面是代碼和相關(guān)注釋:
# 導(dǎo)入庫
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']
# 執(zhí)行推理
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 # 獲取檢測結(jié)果
# 檢查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":
"""
提取用于分類的感興趣區(qū)域(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 # 比例因子,可根據(jù)需要調(diào)整
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)
# 顯示結(jié)果圖像
plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
plt.axis("off")
plt.show()