detector.py

import numpy as np
import torch

from models.yolo.utils.datasets import letterbox
from models.yolo.utils.general import check_img_size, non_max_suppression, scale_coords
from models.yolo.utils.torch_utils import select_device
from models.yolo.model import Model
from models.yolo.experimental import attempt_load


class Detector:

    def __init__(
        self,
        model_file,
        conf_thres = 0.25, 
        iou_thres = 0.45, 
        device = "cpu"
    ):
        self.device = select_device(device)
        weights = torch.load(model_file, self.device)
        self.model = weights["model"]
        self.stride = int(self.model.stride.max())
        self.imgsz = check_img_size(640, s=self.stride)

        self.conf_thres = conf_thres
        self.iou_thres = iou_thres
        
        self.half = device != "cpu"
        if self.half:
            self.model.half()
        self.model.eval()

    def __call__(
        self, 
        img
    ):
        shape = img.shape
        img = letterbox(img, self.imgsz, stride=self.stride)[0]   # resize & padding
        img = img[:, :, ::-1].transpose(2, 0, 1)    # hwc(bgr) -> hwc(rgb) -> c(rgb)hw
        img = np.ascontiguousarray(img)

        img = torch.from_numpy(img).to(self.device)
        img = img.half() if self.half else img.float()
        img /= 255.0

        if img.ndimension() == 3:
            img = img.unsqueeze(0) 

        with torch.no_grad():
            pred = self.model(img, augment=True)[0]

        pred = non_max_suppression(pred, self.conf_thres, self.iou_thres)[0]

        if not len(pred):
            return
        
        # Scale boxes size back
        pred[:, :4] = scale_coords(img.shape[2:], pred[:, :4], shape).round()

        pred_ = pred[:, [0, 1, 2, 3, 5]].to("cpu", int).numpy()
        cls = np.unique(pred_[:, -1], return_index=False)

        if 0 in cls:
            code = "ssq"
        elif 2 in cls:
            code = "cjdlt"
        else:
            return
        
        if not 1 in cls and 3 in cls:
            return

        numbers = pred_[pred_[:, -1] == 1]
        
        issue = pred_[pred_[:, -1] == 3]
        
        return code, issue, numbers