My_Model/object_detection.py

import torch
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import cv2
import os
from My_Model.utilities import get_path, show_image, show_image_with_matplotlib
import transformers

class ObjectDetector:
    def __init__(self):
        self.model = None
        self.processor = None
        self.model_name = None

    def load_model(self, model_name='detic', pretrained=True, model_version='yolov5s'):
        """
        Load the specified object detection model.
        :param model_name: Name of the model to load.
        :param pretrained: Boolean indicating if pretrained model should be used.
        :param model_version: Version of the model, applicable for YOLOv5.
        """
        self.model_name = model_name
        if model_name == 'detic':
            self.load_detic_model(pretrained)
        elif model_name == 'yolov5':
            self.load_yolov5_model(pretrained, model_version)
        else:
            raise ValueError("Unsupported model name")


    def load_detic_model(self, pretrained):
        """Load the Detic model."""
        try:
            model_path = get_path('deformable-detr-detic', 'Models')
            from transformers import AutoImageProcessor, AutoModelForObjectDetection
            self.processor = AutoImageProcessor.from_pretrained(model_path)
            self.model = AutoModelForObjectDetection.from_pretrained(model_path)
        except Exception as e:
            print(f"Error loading Detic model: {e}")


    def load_yolov5_model(self, pretrained, model_version):
        """Load the YOLOv5 model."""
        try:
            model_path = get_path('yolov5', 'Models')
            if model_path and os.path.exists(model_path):
                with os.scandir(model_path) as main_dir:
                    self.model = torch.hub.load(model_path, model_version, pretrained=pretrained, source="local")
            else:
                self.model = torch.hub.load('ultralytics/yolov5', model_version, pretrained=pretrained)
        except Exception as e:
            print(f"Error loading YOLOv5 model: {e}")


    def process_image(self, image_path: str) -> Image.Image:
        """
        Process the image from the given path.
        :param image_path: Path to the image file.
        :return: Processed image.
        """
        with Image.open(image_path) as image:
            return image.convert("RGB")


    def detect_objects(self, image: Image.Image, threshold: float = 0.4):
        """
        Detect objects in the given image.
        :param image: Image in which to detect objects.
        :param threshold: Detection threshold.
        :return: Tuple of detected objects string and list.
        """
        detected_objects_str, detected_objects_list = "", []
        if self.model_name == 'detic':
            detected_objects_str, detected_objects_list = self.detect_with_detic(image, threshold)
        elif self.model_name == 'yolov5':
            detected_objects_str, detected_objects_list = self.detect_with_yolov5(image, threshold)
        return detected_objects_str.strip(), detected_objects_list


    def detect_with_detic(self, image: Image.Image, threshold: float):
        """Detect objects using Detic model."""
        inputs = self.processor(images=image, return_tensors="pt")
        outputs = self.model(**inputs)
        target_sizes = torch.tensor([image.size[::-1]])
        results = self.processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=threshold)[
            0]

        detected_objects_str = ""
        detected_objects_list = []
        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
            if score >= threshold:
                label_name = self.model.config.id2label[label.item()]
                box_rounded = [round(coord, 2) for coord in box.tolist()]
                certainty = round(score.item() * 100, 2)
                detected_objects_str += f"{{object: {label_name}, bounding box: {box_rounded}, certainty: {certainty}%}}\n"
                detected_objects_list.append((label_name, box_rounded, certainty))
        return detected_objects_str, detected_objects_list


    def detect_with_yolov5(self, image: Image.Image, threshold: float):
        """Detect objects using YOLOv5 model."""

        cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
        results = self.model(cv2_img)

        detected_objects_str = ""
        detected_objects_list = []
        for *bbox, conf, cls in results.xyxy[0]:
            if conf >= threshold:
                label_name = results.names[int(cls)]
                box_rounded = [round(coord.item(), 2) for coord in bbox]  # Convert each tensor to float and round
                certainty = round(conf.item() * 100, 2)
                detected_objects_str += f"{{object: {label_name}, bounding box: {box_rounded}, certainty: {certainty}%}}\n"
                detected_objects_list.append((label_name, box_rounded, certainty))
        return detected_objects_str, detected_objects_list


    def draw_boxes(self, image: Image.Image, detected_objects: list, show_confidence: bool = True) -> Image.Image:
        """
        Draw bounding boxes around detected objects in the image.
        :param image: Image on which to draw.
        :param detected_objects: List of detected objects.
        :param show_confidence: Boolean to show confidence scores.
        :return: Image with drawn boxes.
        """
        draw = ImageDraw.Draw(image)
        try:
            font = ImageFont.truetype("arial.ttf", 15)
        except IOError:
            font = ImageFont.load_default()

        colors = ["red", "green", "blue", "yellow", "purple", "orange"]
        label_color_map = {}

        for label_name, box, score in detected_objects:
            if label_name not in label_color_map:
                label_color_map[label_name] = colors[len(label_color_map) % len(colors)]

            color = label_color_map[label_name]
            draw.rectangle(box, outline=color, width=3)

            label_text = f"{label_name}"
            if show_confidence:
                label_text += f" ({round(score, 2)}%)"
            draw.text((box[0], box[1]), label_text, fill=color, font=font)

        return image


if __name__=="__main__":

    detector = ObjectDetector()
    image_path = get_path('horse.jpg', 'Sample_Images')

    detector.load_model('yolov5')  # pass either 'detic' or 'yolov5'

    image = detector.process_image(image_path)
    detected_objects_string, detected_objects_list = detector.detect_objects(image, threshold=0.2)
    image_with_boxes = detector.draw_boxes(image, detected_objects_list, show_confidence=False)
    print(detected_objects_string)
    show_image(image_with_boxes)
    #show_image_with_matplotlib(image_path)