xiongyi/AutoControlSystem-git
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

新增pose模型，注意参数设置

Browse Source
This commit is contained in:
HJW
2025-03-27 10:21:20 +08:00
parent 69e88c2b6b
commit 92778fa76a
6 changed files with 1261 additions and 483 deletions
Download Patch File Download Diff File Expand all files Collapse all files

287
Vision/camera_coordinate_dete.py
View File

@ -19,65 +19,71 @@ from Vision.tool.CameraPe_color2depth import camera_pe as camera_pe_color2depth
from Vision.tool.CameraPe_depth2color import camera_pe as camera_pe_depth2color
from Vision.yolo.yolov8_pt_seg import yolov8_segment
from Vision.yolo.yolov8_openvino import yolov8_segment_openvino
from Vision.yolo.yolov8_pt_pose import yolov8_pose
from Vision.tool.utils import find_position
from Vision.tool.utils import class_names
from Vision.tool.utils import get_disk_space
from Vision.tool.utils import remove_nan_mean_value
from Vision.tool.utils import out_bounds_dete
from Vision.tool.utils import uv_to_XY
from Vision.tool.utils import out_bounds_dete, find_closest_point_index
from Vision.tool.utils import uv_to_XY, shrink_quadrilateral
class Detection:
def __init__(self, use_openvino_model=False, cameraType = 'Pe', alignmentType = 'color2depth'): # cameraType = 'RVC' or cameraType = 'Pe'
def __init__(self, use_openvino_model=False, use_pose_model=True, use_seg_pt_model=True, cameraType = 'Pe', alignmentType = 'color2depth'): # cameraType = 'RVC' or cameraType = 'Pe'
"""
初始化相机及模型
:param use_openvino_model: 加载分割模型
:param use_pose_model: 加载关键点pt模型
:param use_seg_pt_model: 加载分割pt模型
:param use_openvino_model: 选择模型，默认使用openvino
:param cameraType: 选择相机 如本相机 'RVC', 图漾相机 'Pe'
:param alignmentType: 相机对齐方式 color2depth：彩色图对齐深度图 ；depth2color：深度图对齐彩色图
"""
if use_seg_pt_model: # 优先使用pt模型
use_openvino_model = False
elif use_openvino_model:
use_seg_pt_model = False
self.use_openvino_model = use_openvino_model
self.cameraType = cameraType
self.alignmentType= alignmentType
if self.use_openvino_model == False:
model_path = ''.join([os.getcwd(), '/Vision/model/pt/one_bag.pt'])
device = 'cpu'
if self.cameraType == 'RVC':
self.camera_rvc = camera_rvc()
self.seg_distance_threshold = 10 # 1厘米
elif self.cameraType == 'Pe':
if self.alignmentType=='color2depth':
self.camera_rvc = camera_pe_color2depth()
else:
self.camera_rvc = camera_pe_depth2color()
self.seg_distance_threshold = 15 # 2厘米
self.use_pose_model = use_pose_model
self.use_seg_pt_model = use_seg_pt_model
self.alignmentType = alignmentType
if self.cameraType == 'RVC':
self.camera_rvc = camera_rvc()
self.seg_distance_threshold = 10 # 1厘米
elif self.cameraType == 'Pe':
if self.alignmentType == 'color2depth':
self.camera_rvc = camera_pe_color2depth()
else:
print('相机参数错误')
return
self.model = yolov8_segment()
self.model.load_model(model_path, device)
self.camera_rvc = camera_pe_depth2color()
self.seg_distance_threshold = 15 # 2厘米
else:
print('相机参数错误')
return
# 加载openvino-seg
if self.use_openvino_model:
model_path = ''.join([os.getcwd(), '/Vision/model/openvino/one_bag.xml'])
device = 'CPU'
if self.cameraType == 'RVC':
self.camera_rvc = camera_rvc()
self.seg_distance_threshold = 10
elif self.cameraType == 'Pe':
if self.alignmentType == 'color2depth':
self.camera_rvc = camera_pe_color2depth()
else:
self.camera_rvc = camera_pe_depth2color()
self.seg_distance_threshold = 20
else:
print('相机参数错误')
return
self.model = yolov8_segment_openvino(model_path, device, conf_thres=0.3, iou_thres=0.3)
self.model_seg = yolov8_segment_openvino(model_path, device, conf_thres=0.6, iou_thres=0.6)
# 加载pt-seg
if self.use_seg_pt_model:
model_path = ''.join([os.getcwd(), '/Vision/model/pt/one_bag.pt'])
device = 'cpu'
self.model_seg = yolov8_segment()
self.model_seg.load_model(model_path, device)
# 加载pt-pose
if self.use_pose_model:
model_path = ''.join([os.getcwd(), '/Vision/model/pt/one_bag_pose.pt'])
device = 'cpu'
self.model_pose = yolov8_pose(model_path, device)
def get_position(self, Point_isVision=False, Box_isPoint=True, First_Depth =True, Iter_Max_Pixel = 30, save_img_point=0, Height_reduce = 80, width_reduce = 60, Xmin =160, Xmax = 1050, Ymin =290 ,Ymax = 780):
def get_position(self, Use_Pose_Model_Pro=False, Point_isVision=False, Box_isPoint=True, First_Depth =True, Iter_Max_Pixel = 30, save_img_point=0, Height_reduce = 80, width_reduce = 60, Xmin =160, Xmax = 1050, Ymin =290 ,Ymax = 780):
"""
检测料袋相关信息
:param Use_Pose_Model_Pro: True: 选用关键点推理 False : 选用分割模型推理
:param Point_isVision: 点云可视化
:param Box_isPoint: True 返回点云值； False 返回box相机坐标
:param First_Depth: True 返回料袋中心点深度最小的点云值； False 返回面积最大的料袋中心点云值
@ -124,10 +130,19 @@ class Detection:
Abnormal_data_point = point_new.copy()
else:
np.savetxt(save_point_name, point_new)
if self.use_openvino_model == False:
flag, det_cpu, dst_img, masks, category_names = self.model.model_inference(img, 0)
if self.use_pose_model and Use_Pose_Model_Pro:
real_model_pro_isPose = True
else:
flag, det_cpu, scores, masks, category_names = self.model.segment_objects(img)
real_model_pro_isPose = False
if real_model_pro_isPose:
flag, det_cpu, category_names, score_list = self.model_pose.model_inference(img)
else:
if self.use_openvino_model == False:
flag, det_cpu, dst_img, masks, category_names = self.model_seg.model_inference(img, 0)
else:
flag, det_cpu, scores, masks, category_names = self.model_seg.segment_objects(img)
if flag == 1:
xyz = []
nx_ny_nz = []
@ -148,13 +163,23 @@ class Detection:
for i, item in enumerate(det_cpu):
# 画box
box_x1, box_y1, box_x2, box_y2 = item[0:4].astype(np.int32)
if self.use_openvino_model == False:
label = category_names[int(item[5])]
if real_model_pro_isPose:
label = category_names[i]
score = score_list[i]
box_x1 = item[0][0]
box_y1 = item[0][1]
box_x2 = item[3][0]
box_y2 = item[3][1]
pass
else:
label = class_names[int(item[4])]
box_x1, box_y1, box_x2, box_y2 = item[0:4].astype(np.int32)
if self.use_openvino_model == False:
label = category_names[int(item[5])]
score = item[4]
else:
label = class_names[int(item[4])]
score = item[4]
rand_color = (0, 255, 255)
score = item[4]
org = (int((box_x1 + box_x2) / 2), int((box_y1 + box_y2) / 2))
x_center = int((box_x1 + box_x2) / 2)
y_center = int((box_y1 + box_y2) / 2)
@ -164,75 +189,117 @@ class Detection:
thickness=2)
# 画mask
# mask = masks[i].cpu().numpy().astype(int)
if self.use_openvino_model == False:
mask = masks[i].cpu().data.numpy().astype(int)
if real_model_pro_isPose:
# 创建一个与输入数组相同形状的掩码，初始值全为 0
mask = np.zeros(pm.shape[:2], dtype=np.uint8)
# 将四点坐标转换为 numpy 数组
if item[0][0] < item[1][0]:
arr = [[item[0][0], item[0][1]],
[item[1][0], item[1][1]],
[item[3][0], item[3][1]],
[item[2][0], item[2][1]]]
# new_points.reshape((-1, 1, 2))
else:
arr = [[item[3][0], item[3][1]],
[item[2][0], item[2][1]],
[item[0][0], item[0][1]],
[item[1][0], item[1][1]]]
box = arr.copy()
box_outside = arr.copy()
box = shrink_quadrilateral(box, Height_reduce)
pts = np.array(box, np.int32)
# 将四点构成的四边形区域在掩码上标记为 255
cv2.fillPoly(mask, [pts], 255)
# 根据掩码提取对应区域的数据
pm_seg = pm[mask == 255]
# box =[[[item[0][0]+width_reduce, item[0][1]+Height_reduce]],
# [[item[1][0]-width_reduce, item[1][1]+Height_reduce]],
# [[item[3][0]-width_reduce, item[3][1]-Height_reduce]],
# [[item[2][0]+width_reduce, item[2][1]-Height_reduce]]]
box = box.reshape((-1, 1, 2))
# box = np.array(box)
# 内缩
# box_outside = [[[item[0][0], item[0][1]]],
# [[item[1][0], item[1][1]]],
# [[item[3][0], item[3][1]]],
# [[item[2][0], item[2][1]]]] # 外框
box_outside = np.array(box_outside)
box_outside = box_outside.reshape((-1, 1, 2))
# box_outside = np.array(box_outside)
else:
mask = masks[i].astype(int)
mask = mask[box_y1:box_y2, box_x1:box_x2]
if self.use_openvino_model == False:
mask = masks[i].cpu().data.numpy().astype(int)
else:
mask = masks[i].astype(int)
mask = mask[box_y1:box_y2, box_x1:box_x2]
# mask = masks[i].numpy().astype(int)
h, w = box_y2 - box_y1, box_x2 - box_x1
mask_colored = np.zeros((h, w, 3), dtype=np.uint8)
mask_colored[np.where(mask)] = rand_color
##################################
imgray = cv2.cvtColor(mask_colored, cv2.COLOR_BGR2GRAY)
# cv2.imshow('mask',imgray)
# cv2.waitKey(1)
# 2、二进制图像
ret, binary = cv2.threshold(imgray, 10, 255, 0)
# 阈值 二进制图像
# cv2.imshow('bin',binary)
# cv2.waitKey(1)
contours, hierarchy = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
# all_point_list = contours_in(contours)
# print(len(all_point_list))
max_contour = None
max_perimeter = 0
for contour in contours: # 排除小分割区域或干扰区域
perimeter = cv2.arcLength(contour, True)
if perimeter > max_perimeter:
max_perimeter = perimeter
max_contour = contour
# mask = masks[i].numpy().astype(int)
h, w = box_y2 - box_y1, box_x2 - box_x1
mask_colored = np.zeros((h, w, 3), dtype=np.uint8)
mask_colored[np.where(mask)] = rand_color
##################################
imgray = cv2.cvtColor(mask_colored, cv2.COLOR_BGR2GRAY)
# cv2.imshow('mask',imgray)
# cv2.waitKey(1)
# 2、二进制图像
ret, binary = cv2.threshold(imgray, 10, 255, 0)
# 阈值 二进制图像
# cv2.imshow('bin',binary)
# cv2.waitKey(1)
contours, hierarchy = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
# all_point_list = contours_in(contours)
# print(len(all_point_list))
max_contour = None
max_perimeter = 0
for contour in contours: # 排除小分割区域或干扰区域
perimeter = cv2.arcLength(contour, True)
if perimeter > max_perimeter:
max_perimeter = perimeter
max_contour = contour
'''
拟合最小外接矩形，计算矩形中心
'''
'''
拟合最小外接矩形，计算矩形中心
'''
rect = cv2.minAreaRect(max_contour)
if rect[1][0]-width_reduce > 30 and rect[1][1]-Height_reduce > 30:
rect_reduce = (
(rect[0][0], rect[0][1]), (rect[1][0] - width_reduce, rect[1][1] - Height_reduce), rect[2])
else:
rect_reduce = (
(rect[0][0], rect[0][1]), (rect[1][0], rect[1][1]), rect[2])
rect = cv2.minAreaRect(max_contour)
if rect[1][0]-width_reduce > 30 and rect[1][1]-Height_reduce > 30:
rect_reduce = (
(rect[0][0], rect[0][1]), (rect[1][0] - width_reduce, rect[1][1] - Height_reduce), rect[2])
else:
rect_reduce = (
(rect[0][0], rect[0][1]), (rect[1][0], rect[1][1]), rect[2])
# cv2.boxPoints可以将轮廓点转换为四个角点坐标
box_outside = cv2.boxPoints(rect)
# 这一步不影响后面的画图，但是可以保证四个角点坐标为顺时针
startidx = box_outside.sum(axis=1).argmin()
box_outside = np.roll(box_outside, 4 - startidx, 0)
box_outside = np.intp(box_outside)
box_outside = box_outside.reshape((-1, 1, 2)).astype(np.int32)
# cv2.boxPoints可以将轮廓点转换为四个角点坐标
box_outside = cv2.boxPoints(rect)
# 这一步不影响后面的画图，但是可以保证四个角点坐标为顺时针
startidx = box_outside.sum(axis=1).argmin()
box_outside = np.roll(box_outside, 4 - startidx, 0)
box_outside = np.intp(box_outside)
box_outside = box_outside.reshape((-1, 1, 2)).astype(np.int32)
# cv2.boxPoints可以将轮廓点转换为四个角点坐标
box_reduce = cv2.boxPoints(rect_reduce)
startidx = box_reduce.sum(axis=1).argmin()
box_reduce = np.roll(box_reduce, 4 - startidx, 0)
box_reduce = np.intp(box_reduce)
box_reduce = box_reduce.reshape((-1, 1, 2)).astype(np.int32)
# cv2.boxPoints可以将轮廓点转换为四个角点坐标
box_reduce = cv2.boxPoints(rect_reduce)
startidx = box_reduce.sum(axis=1).argmin()
box_reduce = np.roll(box_reduce, 4 - startidx, 0)
box_reduce = np.intp(box_reduce)
box_reduce = box_reduce.reshape((-1, 1, 2)).astype(np.int32)
box_outside = box_outside + [[[box_x1, box_y1]], [[box_x1, box_y1]], [[box_x1, box_y1]],
[[box_x1, box_y1]]]
box = box_reduce + [[[box_x1, box_y1]], [[box_x1, box_y1]], [[box_x1, box_y1]],
[[box_x1, box_y1]]]
'''
提取区域范围内的（x, y）
'''
mask_inside = np.zeros(binary.shape, np.uint8)
cv2.fillPoly(mask_inside, [box_reduce], (255))
pixel_point2 = cv2.findNonZero(mask_inside)
# result = np.zeros_like(color_image)
select_point = []
for i in range(pixel_point2.shape[0]):
select_point.append(pm[pixel_point2[i][0][1]+box_y1, pixel_point2[i][0][0]+box_x1])
select_point = np.array(select_point)
pm_seg = select_point.reshape(-1, 3)
'''
提取区域范围内的（x, y）
'''
mask_inside = np.zeros(binary.shape, np.uint8)
cv2.fillPoly(mask_inside, [box_reduce], (255))
pixel_point2 = cv2.findNonZero(mask_inside)
# result = np.zeros_like(color_image)
select_point = []
for i in range(pixel_point2.shape[0]):
select_point.append(pm[pixel_point2[i][0][1]+box_y1, pixel_point2[i][0][0]+box_x1])
select_point = np.array(select_point)
pm_seg = select_point.reshape(-1, 3)
pm_seg = pm_seg[~np.isnan(pm_seg).all(axis=-1), :] # 剔除 nan
if pm_seg.size < 100:
print("分割点云数量较少，无法拟合平面")
@ -255,9 +322,6 @@ class Detection:
# outlier_cloud.paint_uniform_color([0, 1, 0])
# o3d.visualization.draw_geometries([inlier_cloud, outlier_cloud])
box_outside = box_outside + [[[box_x1, box_y1]], [[box_x1, box_y1]], [[box_x1, box_y1]],[[box_x1, box_y1]]]
box = box_reduce + [[[box_x1, box_y1]], [[box_x1, box_y1]], [[box_x1, box_y1]], [[box_x1, box_y1]]]
box[0][0][1], box[0][0][0] = out_bounds_dete(pm.shape[0], pm.shape[1], box[0][0][1], box[0][0][0])
box[1][0][1], box[1][0][0] = out_bounds_dete(pm.shape[0], pm.shape[1], box[1][0][1], box[1][0][0])
box[2][0][1], box[2][0][0] = out_bounds_dete(pm.shape[0], pm.shape[1], box[2][0][1], box[2][0][0])
@ -277,7 +341,7 @@ class Detection:
point_x, point_y, point_z = remove_nan_mean_value(pm, y_rotation_center, x_rotation_center, iter_max=Iter_Max_Pixel)
if x_rotation_center<Xmin or x_rotation_center>Xmax or y_rotation_center<Ymin or y_rotation_center>Ymax:
continue
cv2.circle(img, (x_rotation_center, y_rotation_center), 4, (255, 255, 255), 5) # 标出中心点
cv2.circle(img, (x_rotation_center, y_rotation_center), 2, (255, 255, 255), 3) # 标出中心点
if np.isnan(point_x): # 点云值为无效值
continue
else:
@ -299,8 +363,11 @@ class Detection:
elif self.cameraType=='Pe':
xyz.append([point_x, point_y, point_z])
Depth_Z.append(point_z)
if real_model_pro_isPose:
RegionalArea.append(0)
else:
RegionalArea.append(cv2.contourArea(max_contour))
nx_ny_nz.append([a, b, c])
RegionalArea.append(cv2.contourArea(max_contour))
uv.append([x_rotation_center, y_rotation_center])
seg_point.append(pm_seg)
cv2.polylines(img, [box], True, (0, 255, 0), 2)
@ -314,7 +381,7 @@ class Detection:
np.savetxt(save_point_name, Abnormal_data_point)
return 1, img, None, None, None
else:
cv2.circle(img, (uv[_idx][0], uv[_idx][1]), 30, (0, 0, 255), 20) # 标出中心点
cv2.circle(img, (uv[_idx][0], uv[_idx][1]), 30, (0, 0, 255), 10) # 标出中心点
if Point_isVision==True:
pcd = o3d.geometry.PointCloud()