How to combine a yolov8 detector and a classifier

[Joeyabuki99]

Search before asking

• I have searched the YOLOv5 issues and discussions and found no similar questions.

Question

Hi guys, I want to combine togheter a detector trained with YOLOv8 and a classifier done with EfficientNetB3.
My detector is been saved with the model.save(output_model_path) and so has a .pt extension, while the classifier is been saved with the same method but has a .h5 extension.

Now how can I combine these two models and test the system on a custom video? I have tried this code but I'm receiving the error TypeError: 'dict' object is not callable on the results = detector(frame) (line 20).
Thanks

` import torch
from tensorflow.keras.models import load_model
import cv2

    detector = torch.load('/myPath/yolov8m_trained.pt')
    
    classifier = load_model('/myPath/efficientnet_model_unfreeze.h5')
    
    video_path = '/myPath//video_test.mp4'
    cap = cv2.VideoCapture(video_path)
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        results = detector(frame)  
        detections = results.xyxy[0].cpu().numpy()  
        
        for detection in detections:
            x1, y1, x2, y2, conf, cls = detection
            
            roi = frame[int(y1):int(y2), int(x1):int(x2)]
            
            roi_resized = cv2.resize(roi, (224, 224))  
            roi_resized = roi_resized / 255.0  
            roi_resized = roi_resized.reshape(1, 224, 224, 3)
    
            pred = classifier.predict(roi_resized)
            class_id = pred.argmax(axis=1)[0]
    
            label = f'Class: {class_id}, Conf: {conf:.2f}'
            cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
            cv2.putText(frame, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
    
        cv2.imshow('Video', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    cap.release()
    cv2.destroyAllWindows()
    
    output_path = 'output_video.mp4'
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, 30.0, (int(cap.get(3)), int(cap.get(4))))
    
    cap = cv2.VideoCapture(video_path)
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
    
        results = detector(frame)
        detections = results.xyxy[0].cpu().numpy()
        
        for detection in detections:
            x1, y1, x2, y2, conf, cls = detection
    
            roi = frame[int(y1):int(y2), int(x1):int(x2)]
    
            roi_resized = cv2.resize(roi, (224, 224))  
            roi_resized = roi_resized / 255.0
            roi_resized = roi_resized.reshape(1, 224, 224, 3)
    
            pred = classifier.predict(roi_resized)
            class_id = pred.argmax(axis=1)[0]
    
            label = f'Class: {class_id}, Conf: {conf:.2f}'
            cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
            cv2.putText(frame, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
    
        out.write(frame)
    
    cap.release()
    out.release()
    cv2.destroyAllWindows()`

[github-actions]

👋 Hello @Joeyabuki99, thank you for your interest in YOLOv5 🚀! Please visit our ⭐️ Tutorials to get started, where you can find quickstart guides for simple tasks like Custom Data Training all the way to advanced concepts like Hyperparameter Evolution.

If this is a 🐛 Bug Report, please provide a minimum reproducible example to help us debug it.

If this is a custom training ❓ Question, please provide as much information as possible, including dataset image examples and training logs, and verify you are following our Tips for Best Training Results.

Requirements

Python>=3.8.0 with all requirements.txt installed including PyTorch>=1.8. To get started:

git clone https://github.com/ultralytics/yolov5  # clone
cd yolov5
pip install -r requirements.txt  # install

Environments

YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Notebooks with free GPU:
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Amazon Deep Learning AMI. See AWS Quickstart Guide
• Docker Image. See Docker Quickstart Guide

Status

If this badge is green, all YOLOv5 GitHub Actions Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv5 training, validation, inference, export and benchmarks on macOS, Windows, and Ubuntu every 24 hours and on every commit.

Introducing YOLOv8 🚀

We're excited to announce the launch of our latest state-of-the-art (SOTA) object detection model for 2023 - YOLOv8 🚀!

Designed to be fast, accurate, and easy to use, YOLOv8 is an ideal choice for a wide range of object detection, image segmentation and image classification tasks. With YOLOv8, you'll be able to quickly and accurately detect objects in real-time, streamline your workflows, and achieve new levels of accuracy in your projects.

Check out our YOLOv8 Docs for details and get started with:

pip install ultralytics

[glenn-jocher]

Hi there! Thanks for reaching out with your question. It looks like you're trying to combine a YOLOv8 detector with an EfficientNetB3 classifier, which is a great approach for leveraging the strengths of both models. Let's address the issue you're encountering.

The error TypeError: 'dict' object is not callable suggests that the detector object might not be correctly instantiated or used. YOLOv8 models typically return a dictionary-like object, and you might need to access specific attributes or methods to get the results.

Here's a revised version of your code with some adjustments:

1. Ensure you are using the latest versions of torch and yolov5 (or yolov8 if applicable).
2. Correctly handle the output from the YOLOv8 model.

import torch
from tensorflow.keras.models import load_model
import cv2

# Load YOLOv8 model
detector = torch.load('/myPath/yolov8m_trained.pt')

# Load EfficientNetB3 classifier
classifier = load_model('/myPath/efficientnet_model_unfreeze.h5')

video_path = '/myPath/video_test.mp4'
cap = cv2.VideoCapture(video_path)

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    
    # Perform detection
    results = detector(frame)
    detections = results['pred'][0].cpu().numpy()  # Adjust based on the actual output structure
    
    for detection in detections:
        x1, y1, x2, y2, conf, cls = detection
        
        roi = frame[int(y1):int(y2), int(x1):int(x2)]
        
        roi_resized = cv2.resize(roi, (224, 224))  
        roi_resized = roi_resized / 255.0  
        roi_resized = roi_resized.reshape(1, 224, 224, 3)

        pred = classifier.predict(roi_resized)
        class_id = pred.argmax(axis=1)[0]

        label = f'Class: {class_id}, Conf: {conf:.2f}'
        cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
        cv2.putText(frame, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

    cv2.imshow('Video', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

output_path = 'output_video.mp4'
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_path, fourcc, 30.0, (int(cap.get(3)), int(cap.get(4))))

cap = cv2.VideoCapture(video_path)
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    results = detector(frame)
    detections = results['pred'][0].cpu().numpy()  # Adjust based on the actual output structure
    
    for detection in detections:
        x1, y1, x2, y2, conf, cls = detection

        roi = frame[int(y1):int(y2), int(x1):int(x2)]

        roi_resized = cv2.resize(roi, (224, 224))  
        roi_resized = roi_resized / 255.0
        roi_resized = roi_resized.reshape(1, 224, 224, 3)

        pred = classifier.predict(roi_resized)
        class_id = pred.argmax(axis=1)[0]

        label = f'Class: {class_id}, Conf: {conf:.2f}'
        cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
        cv2.putText(frame, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

    out.write(frame)

cap.release()
out.release()
cv2.destroyAllWindows()

Please ensure you have the latest versions of torch and yolov5 (or yolov8 if applicable). If the issue persists, providing a minimum reproducible example would help us investigate further. You can find more details on creating a reproducible example here.

Feel free to reach out if you have any more questions or need further assistance. Happy coding! 🚀

[Joeyabuki99]

Hi @glenn-jocher , thank you. I noticed that this problem can be resolved loading the yolo model with detector = YOLO('/myPath/best.pt').
Is this ok or should I load it with the torch model load?

[glenn-jocher]

Hi @Joeyabuki99,

Thank you for your question! It's great to hear that you've found a way to resolve the issue by loading the YOLO model with detector = YOLO('/myPath/best.pt'). This approach is indeed correct and recommended for using YOLO models, as it ensures that all the necessary configurations and dependencies are properly handled.

Using YOLO('/myPath/best.pt') is the preferred method because it leverages the YOLO class's built-in functionalities, which are optimized for loading and running inference with YOLO models. This method abstracts away some of the complexities and potential pitfalls associated with directly using torch.load.

Here's a quick example to illustrate the correct usage:

from yolov5 import YOLO
from tensorflow.keras.models import load_model
import cv2

# Load YOLOv5 model
detector = YOLO('/myPath/best.pt')

# Load EfficientNetB3 classifier
classifier = load_model('/myPath/efficientnet_model_unfreeze.h5')

video_path = '/myPath/video_test.mp4'
cap = cv2.VideoCapture(video_path)

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    
    # Perform detection
    results = detector(frame)
    detections = results.xyxy[0].cpu().numpy()  # Adjust based on the actual output structure
    
    for detection in detections:
        x1, y1, x2, y2, conf, cls = detection
        
        roi = frame[int(y1):int(y2), int(x1):int(x2)]
        
        roi_resized = cv2.resize(roi, (224, 224))  
        roi_resized = roi_resized / 255.0  
        roi_resized = roi_resized.reshape(1, 224, 224, 3)

        pred = classifier.predict(roi_resized)
        class_id = pred.argmax(axis=1)[0]

        label = f'Class: {class_id}, Conf: {conf:.2f}'
        cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
        cv2.putText(frame, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

    cv2.imshow('Video', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

This should work seamlessly for your use case. If you encounter any further issues or have additional questions, please don't hesitate to ask. The YOLO community and the Ultralytics team are here to help! 😊