使用YOLO和Tensorflow模型同时执行对象检测和图像分类

2024年12月26日 由 alex 发表 65 0

对象检测和图像分类是两项不同的任务,各自具有特定的目的。在本文中,我将向您解释对象检测和图像分类是什么,如何训练模型,并在最后同时使用对象检测和图像分类模型对狗的品种进行分类和检测。


1


对象检测

对象检测是一种基本的计算机视觉任务,用于检测和定位物体。简而言之,对象检测模型接受图像作为输入,并输出坐标和标签。


2


如果你需要知道物体的坐标,那么你需要使用对象检测模型,因为图像分类模型不会输出坐标,它们只返回标签。


存在不同的对象检测模型,而且每年训练对象检测模型都变得越来越容易。借助用户友好的库,你可以比自己想象的更容易地训练出你自己的定制模型,并用它来进行检测。


图像分类

图像分类仅输出标签。你给模型输入一张图像,它就会输出一个标签。它更适合对同类物体进行分类。例如,如果你想对海洋动物进行分类,你需要训练一个图像分类模型。


在本文的后面部分,我将与你分享我如何训练一个图像分类模型来对狗的品种进行分类。


3


为什么不只使用对象检测模型?

你可能已经注意到,对象检测模型同时提供坐标和标签作为输出,你可能会想,为什么不干脆对所有任务都使用对象检测模型呢?毕竟,它们在理论上可以同时给出坐标和标签,因此似乎没有必要使用分类模型。你一开始可能会这么想,但有一些你可能没有意识到的不同因素:

  • 对象检测模型在识别和定位场景中的各种物体方面表现出色。但是,在区分外观几乎相同的物体时,图像分类模型通常表现更好(一般来说,并非总是如此)。
  • 你并不总是能找到合适的数据集,而创建数据集可能既耗时又枯燥。如果你决定创建自己的数据集,那么创建图像分类数据集要比创建对象检测数据集容易得多。


4


工作流程

查看下面的图像,它准确地解释了我们现在要做的事情。首先,我们将使用YOLOv8对象检测模型来检测物体,然后从检测到的物体中,我们将尝试使用图像分类模型对它们进行分类。请注意,图像分类模型仅对检测到的物体执行操作,而不是对整个图像执行操作。


5


用于检测狗的对象检测模型

我将使用一个预训练的YOLOv8模型,因为它包含了狗这一类,我将直接使用这个预训练模型。我将使用YOLO模型进行检测,如果检测到狗,我将继续使用图像分类模型。


请记住,一般来说,最好使用针对特定任务的特定数据集来训练模型。


用于狗品种分类的图像分类模型

我将使用TensorFlow来训练一个图像分类模型。根据数据集和参数的不同,训练模型可能需要一些时间,因此,我为狗品种分类模型创建了一个笔记本。


6


结合对象检测 + 图像分类模型

如我之前所解释的,这个过程非常简单。首先执行对象检测模型,然后执行图像分类模型。我尽量为每一行代码都添加了注释来解释,希望一切都清晰明了。


# libraires
import cv2
import numpy as np
from ultralytics import YOLO
from tensorflow.keras.models import load_model
import matplotlib.pyplot as plt
# Load YOLO detection model
yolo_model = YOLO("yolov8s.pt")  # Replace with your YOLO model path
# Load classification model, you can run notebook and save model and use it (check step 2)
classification_model = load_model('dog_classification_model.h5')
# Classification labels
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']


"""
Function to preprocess classification input:
  Before using the classification model, the image needs to be processed. 
  Resizing, normalizing, and adding dimensions are general steps.
  Each model expects a fixed image size, and it is decided before training.
  Here, I trained my model with 180x180 images,
  which is why in the preprocess_image function, I am resizing it to 180x180. 
"""
def preprocess_image(image, target_size):
    img = cv2.resize(image, target_size)  # Resize to target size
    img = img.astype('float32') / 255.0  # Normalize pixel values
    img = np.expand_dims(img, axis=0)  # Add batch dimension
    return img
# Perform inference
def classify_region(image, model, target_size=(180, 180)):  # Size must match the classification model's input
    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
# Load the image
image_path = r"test-images/dog12.jpg"  # Path to your image
image = cv2.imread(image_path)
# YOLO inference --> Object Detection Model
results = yolo_model(image)
detections = results[0].boxes  # Get detections
# Check YOLO label for "dog" and process the bounding boxes
for detection in detections:
    x1, y1, x2, y2 = map(int, detection.xyxy[0].tolist())  # Get bbox coordinates
    conf = float(detection.conf[0])  # Get confidence
    cls_label = yolo_model.names[int(detection.cls[0])]  # Get the label name directly from YOLO
    # Check if the label is "dog"
    if cls_label == "dog":
        """
        Extract the region of interest for classification.
        Remember, the image classification model will perform 
        only on the detected objects, not on the entire image.
        """
        roi = image[y1:y2, x1:x2]
        # Classify the ROI if it's large enough
        if roi.shape[0] > 0 and roi.shape[1] > 0:
            # Image Classification Model
            label = classify_region(roi, classification_model)
            bbox_height = y2 - y1
            font_scale = bbox_height / 200.0  # Scale factor, adjust as needed
            font_thickness = max(1, int(bbox_height / 100))  # Ensure thickness is at least 1
            # Draw the bounding box and label
            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"Detected dog breed: {label}")
cv2.imwrite("dog2-result.jpg",image)
# Display the resulting image
plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
plt.axis("off")
plt.show()


7




文章来源:https://medium.com/gopenai/perform-object-detection-and-image-classification-simultaneously-with-yolo-and-tensorflow-models-8d013f48f114
欢迎关注ATYUN官方公众号
商务合作及内容投稿请联系邮箱:bd@atyun.com
评论 登录
热门职位
Maluuba
20000~40000/月
Cisco
25000~30000/月 深圳市
PilotAILabs
30000~60000/年 深圳市
写评论取消
回复取消