UPDATE Vision 新增行人检测

Vision/detect_person.py

@@ -0,0 +1,123 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+'''
+# @Time    : 2024/9/25 11:47
+# @Author  : hjw
+# @File    : detect_person.py
+'''
+
+import os.path
+import random
+import cv2
+import numpy as np
+import torch
+from ultralytics.nn.autobackend import AutoBackend
+from ultralytics.utils import ops
+
+
+class DetectionPerson:
+    """
+    :param
+    api: None
+    :return: person [bool], img [ndarray]
+    """
+    def __init__(self) -> None:
+        model_path = ''.join([os.getcwd(), '/Vision/model/pt/person.pt'])
+        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_person(self, img_src):
+        person = False
+        # 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]]=="person":
+                    person = True
+                    conf = round(result[4], 2)
+                    print(conf)
+                    self.draw_box(img_src, result[:4], conf, self.names[result[5]], lw, sf, tf)
+
+        return person, 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

get_position_test.py

@@ -9,30 +9,46 @@ from Vision.camera_coordinate_dete import Detection
 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
 import platform
 import cv2
+import os
+#
+"""
+   测试 Vision
+"""
+def detectionPosition_test():
+    detection = Detection()
 
-detection = Detection()
-
-while True:
-    ret, img, xyz, nx_ny_nz, box = detection.get_position()
-    if ret==1:
-        print('xyz点云坐标：', xyz)
-        print('nx_ny_nz法向量：', nx_ny_nz)
-        print('矩形框四顶点：', box)
-        # img = cv2.resize(img,(720, 540))
-        if xyz!=None:
-            transformation_matrix = R_matrix(521.534, 0.705, 850.03, 0.0, 90.0, 0.0)   # (x, y, z, u, v, w)
-            target_position, noraml_base = getPosition(xyz[0], xyz[1], xyz[2], nx_ny_nz[0], nx_ny_nz[1], nx_ny_nz[2], transformation_matrix, box)
-            print("target_position:", target_position)
-            print("noraml_base", noraml_base)
-        cv2.imshow('img', img)
-        cv2.waitKey(1)
+    while True:
+        ret, img, xyz, nx_ny_nz, box = detection.get_position()
+        if ret==1:
+            print('xyz点云坐标：', xyz)
+            print('nx_ny_nz法向量：', nx_ny_nz)
+            print('矩形框四顶点：', box)
+            # img = cv2.resize(img,(720, 540))
+            if xyz!=None:
+                transformation_matrix = R_matrix(521.534, 0.705, 850.03, 0.0, 90.0, 0.0)   # (x, y, z, u, v, w)
+                target_position, noraml_base = getPosition(xyz[0], xyz[1], xyz[2], nx_ny_nz[0], nx_ny_nz[1], nx_ny_nz[2], transformation_matrix, box)
+                print("target_position:", target_position)
+                print("noraml_base", noraml_base)
+            cv2.imshow('img', img)
+            cv2.waitKey(1)
 
 
-# point = np.loadtxt('./Vision/model/data/test_nan.xyz').reshape((6, 8, 3))
-# a = point.copy()
-# c = a.tolist()
-# x, y, z = remove_nan_mean_value(point,3,3)
-# print(x, y, z )
-# pass
+def detectionPerson_test():
+    detectionPerson = DetectionPerson()
+    video_path = ''.join([os.getcwd(), '/Vision/model/data/1.mp4'])
+    cam = cv2.VideoCapture(video_path)
+    while cam.isOpened():
+        ret, frame = cam.read()
+        if ret:
+            person, img = detectionPerson.get_person(frame)
+            print(person)
+            cv2.imshow('img', img)
+            cv2.waitKey(1)
+        else:
+            break
+
+if __name__ == '__main__':
+    detectionPerson_test()