UPDATE Vision 新增料袋数量检测

Vision/camera_coordinate_dete.py

@@ -26,7 +26,7 @@ import os
 
 class Detection:
 
-    def __init__(self, use_openvino_model=True, cameraType = 'RVC'):                   # cameraType = 'RVC'  or cameraType = 'Pe'
+    def __init__(self, use_openvino_model=False, cameraType = 'RVC'):                   # cameraType = 'RVC'  or cameraType = 'Pe'
         """
             初始化相机及模型
             :param use_openvino_model:  选择模型，默认使用openvino

Vision/camera_coordinate_dete_img.py

@@ -26,7 +26,7 @@ import os
 
 class Detection:
 
-    def __init__(self, use_openvino_model=True, cameraType = 'RVC'):
+    def __init__(self, use_openvino_model=False, cameraType = 'RVC'):
         self.use_openvino_model = use_openvino_model
         self.cameraType = cameraType
         if self.use_openvino_model == False:

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

Vision/tool/CameraHIK.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+'''
+# @Time    : 2024/10/11 10:43
+# @Author  : hjw
+# @File    : CameraHIK.py
+'''
+import cv2
+import socket
+
+def portisopen(ip, port):
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.settimeout(1)
+    state = sock.connect_ex((ip, port))
+    if 0 == state:
+        # print("port is open")
+        return True
+    else:
+        # print("port is closed")
+        return False
+
+class camera_HIK():
+
+    def __init__(self, ip, port, name, pw):
+        # "rtsp://admin:zlzk.123@192.168.1.64:554"
+        self.camera_url = "rtsp://" + name + ":" + pw + "@" + ip + ":" + str(port)
+        ret = portisopen(ip, port)
+        self.ip = ip
+        self.port = port
+        self.init_success = False
+        if ret:
+            self.cap = cv2.VideoCapture(self.camera_url)
+            self.init_sucess = True
+        else:
+            print('海康摄像头网络错误，请检测IP')
+
+    def get_img(self):
+        ret = False
+        frame = None
+        if self.init_success==True:
+            if portisopen(self.ip, self.port):
+                ret, frame = self.cap.read()
+            else:
+                print('海康摄像头网络断开')
+        else:
+            if portisopen(self.ip, self.port):
+                self.reconnect_camera()
+                ret, frame = self.cap.read()
+            else:
+                print('海康摄像头网络断开')
+
+        return ret, frame
+
+    def reconnect_camera(self):
+        if self.init_success == True:
+            self.cap.release()
+        self.cap = cv2.VideoCapture(self.camera_url)
+        print("海康摄像头重连")
+
+    def release_camera(self):
+        self.cap.release()
+
+

get_position_test.py

@@ -10,6 +10,8 @@ from Trace.handeye_calibration import *
 from Vision.tool.utils import get_disk_space
 from Vision.tool.utils import remove_nan_mean_value
 from Vision.detect_person import DetectionPerson
+from Vision.detect_bag_num import DetectionBagNum
+from Vision.tool.CameraHIK import camera_HIK
 import platform
 import cv2
 import os
@@ -61,6 +63,28 @@ def detectionPerson_test():
         else:
             break
 
+def HIK_Camera_test():
+    MY_CAMERA = camera_HIK("192.168.1.121", 554, "admin", "zlzk.123")
+    while True:
+        ret, frame = MY_CAMERA.get_img()
+        if ret:
+            img = cv2.resize(frame, (800, 600))
+            cv2.imshow('img', img)
+            cv2.waitKey(1)
+def detectionBagNum_test():
+    detectionBagNum = DetectionBagNum()
+    video_path = ''.join([os.getcwd(), '/Vision/model/data/2.mp4'])
+    cam = cv2.VideoCapture(video_path)
+    while cam.isOpened():
+        ret, frame = cam.read()
+        if ret:
+            num, img = detectionBagNum.get_BagNum(frame)
+            print(num)
+            cv2.imshow('img', img)
+            cv2.waitKey(1)
+        else:
+            break
+
 
 if __name__ == '__main__':
-    detectionPosition_test()
+    detectionBagNum_test()