controllers/DetecInfoBoxes/GetBoxes.py

import cv2
import torch
import numpy as np

from DetecInfoBoxes.models.experimental import attempt_load
from DetecInfoBoxes.utils.general import check_img_size, non_max_suppression, scale_coords, set_logging
from DetecInfoBoxes.utils.torch_utils import select_device


class Detect:
    def __init__(self, opt):
        self.opt = opt

    def load_model(self, weights):
        """Hàm load model YOLO V7

                Args:
                    weights (str): Đường dẫn đến file trọng số model YOLO V7 (.pt) hoặc (.pth).

                Returns:
                    imgsz, stride, device, half, model, names: Các tham số đầu ra phục vụ cho quá trình predict của model.

                """
        with torch.no_grad():
            imgsz = self.opt['img-size']
            set_logging()
            device = select_device(self.opt['device'])
            half = device.type != 'cpu'
            model = attempt_load(weights, map_location=device)
            stride = int(model.stride.max())
            imgsz = check_img_size(imgsz, s=stride)
            if half:
                model.half()

            names = model.module.names if hasattr(model, 'module') else model.names
            if device.type != 'cpu':
                model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))

        return imgsz, stride, device, half, model, names

    @staticmethod
    def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True,
                  stride=32):

        shape = img.shape[:2]
        if isinstance(new_shape, int):
            new_shape = (new_shape, new_shape)

        r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
        if not scaleup:
            r = min(r, 1.0)

        ratio = r, r
        new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
        dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]
        if auto:
            dw, dh = np.mod(dw, stride), np.mod(dh, stride)
        elif scaleFill:
            dw, dh = 0.0, 0.0
            new_unpad = (new_shape[1], new_shape[0])
            ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]

        dw /= 2
        dh /= 2

        if shape[::-1] != new_unpad:  # resize
            img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
        top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
        left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
        img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)

        return img, ratio, (dw, dh)

    @staticmethod
    def dict_processing(dictionary: dict):
        """Hàm chuyển đổi dict từ dạng {label: [x_min, y_min, x_max, y_max, conf]} sang dạng {label: [x_center, y_center, width, height, conf, label]}.

        Args:
            dictionary (dict): dictionary đầu vào.

        Returns:
            new_dictionary (dict): dictionary sau khi được chuyển đổi sang dạng mới (new_dictionary).

        """
        new_dictionary = {}
        for item in dictionary.items():
            new_bbox = []
            for bbox in item[1]:
                new_bbox.append([int((bbox[0] + bbox[2]) / 2), int((bbox[1] + bbox[3]) / 2), int(bbox[2] - bbox[0]),
                                 int(bbox[3] - bbox[1]), bbox[4], item[0]])
            new_dictionary.update({item[0]: new_bbox})
        return new_dictionary

    def prediction(self, img_official: np.ndarray, imgsz, stride, device, half, model, names):
        """Hàm đưa ra dự đoán của model YOLO

        Args:
            img_official (np.ndarray): Ảnh đầu vào cho model dự đoán.
            imgsz, stride, device, half, model, names: Các tham số có được từ quá trình load model.

        Returns:
            dictionary (dict): dictionary dạng {label: [x_min, y_min, x_max, y_max, conf]}.
        """
        img = self.letterbox(img_official, imgsz, stride=stride)[0]
        img = img[:, :, ::-1].transpose(2, 0, 1)
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(device)
        img = img.half() if half else img.float()
        img /= 255.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        pred = model(img, augment=False)[0]

        classes = None
        pred = non_max_suppression(pred, self.opt['conf-thres'], self.opt['iou-thres'], classes=classes,
                                   agnostic=False)
        dictsLabel = {}
        for i, det in enumerate(pred):
            s = ''
            s += '%gx%g ' % img.shape[2:]
            if len(det):
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img_official.shape).round()

                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()
                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "

                for *xyxy, conf, cls in reversed(det):
                    label = f'{names[int(cls)]}'
                    confident = f'{conf:.2f}'

                    if label not in dictsLabel.keys():
                        dictsLabel.update({label: [[xyxy[0], xyxy[1], xyxy[2], xyxy[3], confident]]})
                    else:
                        values = dictsLabel[label]
                        values.append([xyxy[0], xyxy[1], xyxy[2], xyxy[3], confident])
                        dictsLabel.update({label: values})
        return dictsLabel

    # Used for testing
    @staticmethod
    def new_draw_boxes(new_dicts: dict, image_draw: np.ndarray):
        """Hàm vẽ kết quả detect của model YOLO

        Args:
            new_dicts (dict): Dictionary (dạng {label: [x_center, y_center, width, height, conf, label]}) thu được sau khi chuyển
            đổi thông qua hàm dict_processing().
            image_draw (np.ndarray): Ảnh để vẽ dự đoán của model.

        Returns:
            None
        """
        for label in new_dicts:
            horizontal_list = new_dicts.get(label)
            for bbox in horizontal_list:
                x_center = int(bbox[0])
                y_center = int(bbox[1])
                w = int(bbox[2])
                h = int(bbox[3])
                cv2.putText(image_draw, str(bbox[5]),
                            (x_center - int(w / 2) + 1, y_center - int(h / 2) - 10),
                            cv2.FONT_HERSHEY_TRIPLEX, 0.5, (0, 0, 255))
                cv2.rectangle(image_draw, (x_center - int(w / 2), y_center - int(h / 2)),
                              (x_center + int(w / 2), y_center + int(h / 2)), (0, 0, 255), 1)

    @staticmethod
    def draw_boxes(dicts: dict, image_draw: np.ndarray):
        """Hàm vẽ kết quả dự đoán của model YOLO

        Args:
            dicts (dict): Dictionary (dạng {label: [x_min, y_min, x_max, y_max, conf]}) thu được sau khi model dự đoán.
            image_draw (np.ndarray): Ảnh để vẽ dự đoán của model.

        Returns:
            None
        """
        for label in dicts:
            horizontal_list = dicts.get(label)
            for bbox in horizontal_list:
                x_min = int(bbox[0])
                y_min = int(bbox[1])
                x_max = int(bbox[2])
                y_max = int(bbox[3])
                cv2.putText(image_draw, str(label) + ' ' + str(bbox[4]), (x_min + 1, y_min - 10),
                            cv2.FONT_HERSHEY_TRIPLEX, 0.5, (255, 0, 255))
                cv2.rectangle(image_draw, (x_min, y_min),
                              (x_max, y_max), (255, 0, 255), 1)