ailai/Vision/detect_bag_num.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
# @Time    : 2024/10/11 9:44
# @Author  : hjw
# @File    : detect_bag_num.py
'''

import os.path
import random
import cv2
import numpy as np
import torch
from Vision.tool.CameraHIK import camera_HIK
from ultralytics.nn.autobackend import AutoBackend
from ultralytics.utils import ops

class DetectionBagNum:
    """
    检测顶层料袋数量
    :param
    api: camera -> ip , port, name, pw
    :return ret :[bool] 相机是否正常工作
    :return BagNum :[int]  返回料袋数量
    :return img_src :[ndarray]  返回检测图像
    """
    def __init__(self, ip="192.168.1.121", port=554, name="admin", pw="zlzk.123"):
        model_path = ''.join([os.getcwd(), '/Vision/model/pt/bagNum.pt'])
        self.HIk_Camera = camera_HIK(ip, port, name, pw)
        self.imgsz = 640
        self.cuda = 'cpu'
        self.conf = 0.40
        self.iou = 0.45
        self.model = AutoBackend(model_path, device=torch.device(self.cuda))
        self.model.eval()
        self.names = self.model.names
        self.half = False
        self.color = {"font": (255, 255, 255)}
        self.color.update(
            {self.names[i]: (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
             for i in range(len(self.names))})

    def get_BagNum(self):
        """
            检测顶层料袋数量
            :param
            api: camera -> ip , port, name, pw
            :return ret :[bool] 相机是否正常工作
            :return BagNum :[int]  返回料袋数量
            :return img_src :[ndarray]  返回检测图像
        """
        BagNum = 0
        ret, img_src = self.HIk_Camera.get_img()
        if ret:
            # img_src = cv2.imread(img_path)
            img = self.precess_image(img_src, self.imgsz, self.half, self.cuda)
            preds = self.model(img)
            det = ops.non_max_suppression(preds, self.conf, self.iou, classes=None, agnostic=False, max_det=300,
                                          nc=len(self.names))

            for i, pred in enumerate(det):
                lw = max(round(sum(img_src.shape) / 2 * 0.003), 2)  # line width
                tf = max(lw - 1, 1)  # font thickness
                sf = lw / 3  # font scale
                pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], img_src.shape)
                results = pred.cpu().detach().numpy()
                # cv2.imshow('img2', img_src)
                # cv2.waitKey(1)
                for result in results:
                    if self.names[result[5]]=="bag":
                        BagNum += 1
                        conf = round(result[4], 2)
                        self.draw_box(img_src, result[:4], conf, self.names[result[5]], lw, sf, tf)

        return ret, BagNum, img_src

    def draw_box(self, img_src, box, conf, cls_name, lw, sf, tf):
        color = self.color[cls_name]
        conf = str(conf)
        label = f'{cls_name} {conf}'
        p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
        # 绘制矩形框
        cv2.rectangle(img_src, p1, p2, color, thickness=lw, lineType=cv2.LINE_AA)
        # text width, height
        w, h = cv2.getTextSize(label, 0, fontScale=sf, thickness=tf)[0]
        # label fits outside box
        outside = box[1] - h - 3 >= 0
        p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
        # 绘制矩形框填充
        cv2.rectangle(img_src, p1, p2, color, -1, cv2.LINE_AA)
        # 绘制标签
        cv2.putText(img_src, label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2),
                    0, sf, self.color["font"], thickness=2, lineType=cv2.LINE_AA)

    @staticmethod
    def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), scaleup=True, stride=32):
        # Resize and pad image while meeting stride-multiple constraints
        shape = im.shape[:2]  # current shape [height, width]
        if isinstance(new_shape, int):
            new_shape = (new_shape, new_shape)

        # Scale ratio (new / old)
        r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
        if not scaleup:  # only scale down, do not scale up (for better val mAP)
            r = min(r, 1.0)

        # Compute padding
        ratio = r, r  # width, height ratios
        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]  # wh padding
        # minimum rectangle
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
        dw /= 2  # divide padding into 2 sides
        dh /= 2

        if shape[::-1] != new_unpad:  # resize
            im = cv2.resize(im, 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))
        im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
        return im, ratio, (dw, dh)

    def precess_image(self, img_src, img_size, half, device):
        # Padded resize
        img = self.letterbox(img_src, img_size)[0]
        # Convert
        img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(device)

        img = img.half() if half else img.float()  # uint8 to fp16/32
        img = img / 255  # 0 - 255 to 0.0 - 1.0
        if len(img.shape) == 3:
            img = img[None]  # expand for batch dim
        return img