first commit

Vision/detect_bag_num.py

@@ -0,0 +1,136 @@
+#!/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