diff --git a/Vision/camera_coordinate_dete.py b/Vision/camera_coordinate_dete.py index 1864826..50351ab 100644 --- a/Vision/camera_coordinate_dete.py +++ b/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 diff --git a/Vision/camera_coordinate_dete_img.py b/Vision/camera_coordinate_dete_img.py index d740b86..edbc682 100644 --- a/Vision/camera_coordinate_dete_img.py +++ b/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: diff --git a/Vision/detect_bag_num.py b/Vision/detect_bag_num.py new file mode 100644 index 0000000..32d63dc --- /dev/null +++ b/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 \ No newline at end of file diff --git a/Vision/model/pt/bagNum.pt b/Vision/model/pt/bagNum.pt new file mode 100644 index 0000000..645294e Binary files /dev/null and b/Vision/model/pt/bagNum.pt differ diff --git a/Vision/tool/CameraHIK.py b/Vision/tool/CameraHIK.py new file mode 100644 index 0000000..95f824d --- /dev/null +++ b/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() + + diff --git a/get_position_test.py b/get_position_test.py index 43f29df..9e0247a 100644 --- a/get_position_test.py +++ b/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() \ No newline at end of file + detectionBagNum_test() \ No newline at end of file