forked from huangxin/ailai
394 lines
12 KiB
Python
394 lines
12 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: UTF-8 -*-
|
||
'''
|
||
@Project ：AutoControlSystem-master
|
||
@File ：utils.py
|
||
@IDE ：PyCharm
|
||
@Author ：hjw
|
||
@Date ：2024/8/29 15:07
|
||
'''
|
||
|
||
import numpy as np
|
||
import cv2
|
||
import psutil
|
||
from psutil._common import bytes2human
|
||
|
||
|
||
def uv_to_XY(cameraType, u, v):
|
||
"""
|
||
像素坐标转相机坐标
|
||
Args:
|
||
cameraType:
|
||
u:
|
||
v:
|
||
|
||
Returns:
|
||
如本：
|
||
|
||
ExtrinsicMatrix:
|
||
[-0.9916700124740601, -0.003792409785091877, 0.12874870002269745, 0.10222162306308746, -0.003501748666167259, 0.9999907612800598, 0.002483875723555684, -0.08221593499183655, -0.12875692546367645, 0.0020123394206166267, -0.9916741251945496, 0.6480034589767456, 0.0, 0.0, 0.0, 1.0]
|
||
|
||
IntrinsicParameters:
|
||
[2402.101806640625, 0.0, 739.7069091796875,
|
||
0.0, 2401.787353515625, 584.73046875,
|
||
0.0, 0.0, 1.0]
|
||
|
||
distortion:
|
||
[-0.04248141124844551, 0.24386045336723328, -0.38333430886268616, -0.0017840253422036767, 0.0007602088153362274]
|
||
|
||
图漾：
|
||
|
||
depth image format list:
|
||
0 -size[640x480] - desc:DEPTH16_640x480
|
||
1 -size[1280x960] - desc:DEPTH16_1280x960
|
||
2 -size[320x240] - desc:DEPTH16_320x240
|
||
delth calib info:
|
||
calib size :[1280x960]
|
||
calib intr :
|
||
(1048.3614501953125, 0.0, 652.146240234375,
|
||
0.0, 1048.3614501953125, 500.26397705078125,
|
||
0.0, 0.0, 1.0)
|
||
calib extr : (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
|
||
calib distortion : (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
|
||
|
||
"""
|
||
x = None
|
||
y = None
|
||
zc = 1 # 设深度z为1
|
||
if cameraType == 'RVC':
|
||
u0 = 739.70
|
||
v0 = 584.73
|
||
fx = 2402.10
|
||
fy = 2401.78
|
||
x = (u-u0)*zc/fx
|
||
y = (v-v0)*zc/fy
|
||
elif cameraType == 'Pe':
|
||
u0 = 652.14
|
||
v0 = 500.26
|
||
fx = 1048.36
|
||
fy = 1048.36
|
||
|
||
x = (u - u0) * zc / fx
|
||
y = (v - v0) * zc / fy
|
||
return x, y, zc
|
||
|
||
|
||
|
||
def out_bounds_dete(pm_y, pm_x, piont_y, piont_x):
|
||
if piont_y>=pm_y:
|
||
piont_y = pm_y-1
|
||
print('四坐标点超出点云大小')
|
||
if piont_y<0:
|
||
piont_y=0
|
||
print('四坐标点超出点云大小')
|
||
if piont_x>=pm_x:
|
||
piont_x = pm_x-1
|
||
print('四坐标点超出点云大小')
|
||
if piont_x<0:
|
||
piont_x=0
|
||
print('四坐标点超出点云大小')
|
||
return piont_y, piont_x
|
||
|
||
def remove_nan_mean_value(pm, y, x, iter_max=50):
|
||
y, x = out_bounds_dete(pm.shape[0], pm.shape[1], y, x)
|
||
point_x, point_y, point_z = pm[y, x]
|
||
if np.isnan(point_x):
|
||
point_x_list = []
|
||
point_y_list = []
|
||
point_z_list = []
|
||
iter_current = 1
|
||
pm_shape_y = pm.shape[0]
|
||
pm_shape_x = pm.shape[1]
|
||
remove_nan_isok = False
|
||
print('Nan值去除')
|
||
while iter_current < iter_max:
|
||
# 计算开始点
|
||
if y - iter_current > 0:
|
||
y_start = y - iter_current
|
||
else:
|
||
y_start = 0
|
||
|
||
if x - iter_current > 0:
|
||
x_start = x - iter_current
|
||
else:
|
||
x_start = 0
|
||
|
||
for idx_y in range(iter_current*2 + 1):
|
||
y_current = y_start + idx_y
|
||
if y_current > pm_shape_y-1:
|
||
continue
|
||
for idx_x in range(iter_current*2 + 1):
|
||
x_current = x_start + idx_x
|
||
if x_current > pm_shape_x-1:
|
||
continue
|
||
elif np.isnan(pm[y_current, x_current][0]) == False:
|
||
point_x_list.append(pm[y_current, x_current][0])
|
||
point_y_list.append(pm[y_current, x_current][1])
|
||
point_z_list.append(pm[y_current, x_current][2])
|
||
|
||
len_point_x = len(point_x_list)
|
||
if len_point_x > 0:
|
||
point_x = sum(point_x_list)/len_point_x
|
||
point_y = sum(point_y_list)/len_point_x
|
||
point_z = sum(point_z_list)/len_point_x
|
||
remove_nan_isok = True
|
||
break
|
||
iter_current += 1
|
||
else:
|
||
remove_nan_isok = True
|
||
if remove_nan_isok == True:
|
||
return point_x, point_y, point_z
|
||
else:
|
||
print(f'在{iter_max}*{iter_max}范围中未找到有效值，所有点云值为无效值')
|
||
return np.nan, np.nan, np.nan
|
||
|
||
def remove_nan(pm, y, x):
|
||
point_x, point_y, point_z = pm[y, x]
|
||
if np.isnan(point_x):
|
||
for i in range(10):
|
||
point_x, point_y, point_z = pm[y+i, x]
|
||
if np.isnan(point_x)==False:
|
||
break
|
||
return point_x, point_y, point_z
|
||
|
||
|
||
def get_disk_space(path='C:'):
|
||
|
||
usage = psutil.disk_usage(path)
|
||
space_free = bytes2human(usage.free)
|
||
# space_total = bytes2human(usage.total)
|
||
# space_used = bytes2human(usage.used)
|
||
# space_free = bytes2human(usage.free)
|
||
# space_used_percent = bytes2human(usage.percent)
|
||
space_free = float(space_free[:-1])
|
||
return space_free
|
||
def find_position(Depth_Z, RegionalArea, RegionalArea_Threshold, first_depth=True):
|
||
if first_depth == True:
|
||
sorted_id = sorted(range(len(Depth_Z)), key=lambda k: Depth_Z[k], reverse=False)
|
||
# Depth_Z1 = [Depth_Z[i] for i in sorted_id]
|
||
# RegionalArea1 = [RegionalArea[i] for i in sorted_id]
|
||
# for i in range(len(Depth_Z1)):
|
||
# if RegionalArea1[i] > RegionalArea_Threshold:
|
||
# return sorted_id[i]
|
||
if len(sorted_id)>0:
|
||
return sorted_id[0]
|
||
else:
|
||
return
|
||
|
||
else:
|
||
sorted_id = sorted(range(len(RegionalArea)), key=lambda k: RegionalArea[k], reverse=True)
|
||
# Depth_Z1 = [Depth_Z[i] for i in sorted_id]
|
||
# RegionalArea1 = [RegionalArea[i] for i in sorted_id]
|
||
# for i in range(len(Depth_Z1)):
|
||
# if RegionalArea1[i] > RegionalArea_Threshold:
|
||
# return sorted_id[i]
|
||
if len(sorted_id)>0:
|
||
return sorted_id[0]
|
||
else:
|
||
return
|
||
|
||
class_names = ['box', 'other']
|
||
|
||
# Create a list of colors for each class where each color is a tuple of 3 integer values
|
||
rng = np.random.default_rng(3)
|
||
colors = rng.uniform(0, 255, size=(len(class_names), 3))
|
||
|
||
|
||
def nms(boxes, scores, iou_threshold):
|
||
# Sort by score
|
||
sorted_indices = np.argsort(scores)[::-1]
|
||
|
||
keep_boxes = []
|
||
while sorted_indices.size > 0:
|
||
# Pick the last box
|
||
box_id = sorted_indices[0]
|
||
keep_boxes.append(box_id)
|
||
|
||
# Compute IoU of the picked box with the rest
|
||
ious = compute_iou(boxes[box_id, :], boxes[sorted_indices[1:], :])
|
||
|
||
# Remove boxes with IoU over the threshold
|
||
keep_indices = np.where(ious < iou_threshold)[0]
|
||
|
||
# print(keep_indices.shape, sorted_indices.shape)
|
||
sorted_indices = sorted_indices[keep_indices + 1]
|
||
|
||
return keep_boxes
|
||
|
||
|
||
def compute_iou(box, boxes):
|
||
# Compute xmin, ymin, xmax, ymax for both boxes
|
||
xmin = np.maximum(box[0], boxes[:, 0])
|
||
ymin = np.maximum(box[1], boxes[:, 1])
|
||
xmax = np.minimum(box[2], boxes[:, 2])
|
||
ymax = np.minimum(box[3], boxes[:, 3])
|
||
|
||
# Compute intersection area
|
||
intersection_area = np.maximum(0, xmax - xmin) * np.maximum(0, ymax - ymin)
|
||
|
||
# Compute union area
|
||
box_area = (box[2] - box[0]) * (box[3] - box[1])
|
||
boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
|
||
union_area = box_area + boxes_area - intersection_area
|
||
|
||
# Compute IoU
|
||
iou = intersection_area / union_area
|
||
|
||
return iou
|
||
|
||
|
||
def xywh2xyxy(x):
|
||
# Convert bounding box (x, y, w, h) to bounding box (x1, y1, x2, y2)
|
||
y = np.copy(x)
|
||
y[..., 0] = x[..., 0] - x[..., 2] / 2
|
||
y[..., 1] = x[..., 1] - x[..., 3] / 2
|
||
y[..., 2] = x[..., 0] + x[..., 2] / 2
|
||
y[..., 3] = x[..., 1] + x[..., 3] / 2
|
||
return y
|
||
|
||
|
||
def sigmoid(x):
|
||
return 1 / (1 + np.exp(-x))
|
||
|
||
|
||
def draw_detections(image, boxes, scores, class_ids, mask_alpha=0.3, mask_maps=None):
|
||
img_height, img_width = image.shape[:2]
|
||
size = min([img_height, img_width]) * 0.0006
|
||
text_thickness = int(min([img_height, img_width]) * 0.001)
|
||
|
||
mask_img = draw_masks(image, boxes, class_ids, mask_alpha, mask_maps)
|
||
|
||
# Draw bounding boxes and labels of detections
|
||
for box, score, class_id in zip(boxes, scores, class_ids):
|
||
color = colors[class_id]
|
||
|
||
x1, y1, x2, y2 = box.astype(int)
|
||
|
||
# Draw rectangle
|
||
cv2.rectangle(mask_img, (x1, y1), (x2, y2), color, 2)
|
||
|
||
label = class_names[class_id]
|
||
caption = f'{label} {int(score * 100)}%'
|
||
(tw, th), _ = cv2.getTextSize(text=caption, fontFace=cv2.FONT_HERSHEY_SIMPLEX,
|
||
fontScale=size, thickness=text_thickness)
|
||
th = int(th * 1.2)
|
||
|
||
cv2.rectangle(mask_img, (x1, y1),
|
||
(x1 + tw, y1 - th), color, -1)
|
||
|
||
cv2.putText(mask_img, caption, (x1, y1),
|
||
cv2.FONT_HERSHEY_SIMPLEX, size, (255, 255, 255), text_thickness, cv2.LINE_AA)
|
||
|
||
return mask_img
|
||
|
||
|
||
def draw_masks(image, boxes, class_ids, mask_alpha=0.3, mask_maps=None):
|
||
mask_img = image.copy()
|
||
|
||
# Draw bounding boxes and labels of detections
|
||
for i, (box, class_id) in enumerate(zip(boxes, class_ids)):
|
||
color = colors[class_id]
|
||
|
||
x1, y1, x2, y2 = box.astype(int)
|
||
|
||
# Draw fill mask image
|
||
if mask_maps is None:
|
||
cv2.rectangle(mask_img, (x1, y1), (x2, y2), color, -1)
|
||
else:
|
||
crop_mask = mask_maps[i][y1:y2, x1:x2, np.newaxis]
|
||
crop_mask_img = mask_img[y1:y2, x1:x2]
|
||
crop_mask_img = crop_mask_img * (1 - crop_mask) + crop_mask * color
|
||
mask_img[y1:y2, x1:x2] = crop_mask_img
|
||
|
||
return cv2.addWeighted(mask_img, mask_alpha, image, 1 - mask_alpha, 0)
|
||
|
||
|
||
def draw_comparison(img1, img2, name1, name2, fontsize=2.6, text_thickness=3):
|
||
(tw, th), _ = cv2.getTextSize(text=name1, fontFace=cv2.FONT_HERSHEY_DUPLEX,
|
||
fontScale=fontsize, thickness=text_thickness)
|
||
x1 = img1.shape[1] // 3
|
||
y1 = th
|
||
offset = th // 5
|
||
cv2.rectangle(img1, (x1 - offset * 2, y1 + offset),
|
||
(x1 + tw + offset * 2, y1 - th - offset), (0, 115, 255), -1)
|
||
cv2.putText(img1, name1,
|
||
(x1, y1),
|
||
cv2.FONT_HERSHEY_DUPLEX, fontsize,
|
||
(255, 255, 255), text_thickness)
|
||
|
||
(tw, th), _ = cv2.getTextSize(text=name2, fontFace=cv2.FONT_HERSHEY_DUPLEX,
|
||
fontScale=fontsize, thickness=text_thickness)
|
||
x1 = img2.shape[1] // 3
|
||
y1 = th
|
||
offset = th // 5
|
||
cv2.rectangle(img2, (x1 - offset * 2, y1 + offset),
|
||
(x1 + tw + offset * 2, y1 - th - offset), (94, 23, 235), -1)
|
||
|
||
cv2.putText(img2, name2,
|
||
(x1, y1),
|
||
cv2.FONT_HERSHEY_DUPLEX, fontsize,
|
||
(255, 255, 255), text_thickness)
|
||
|
||
combined_img = cv2.hconcat([img1, img2])
|
||
if combined_img.shape[1] > 3840:
|
||
combined_img = cv2.resize(combined_img, (3840, 2160))
|
||
|
||
return combined_img
|
||
|
||
|
||
|
||
def fit_plane_vision(box_list, normal_vector):
|
||
plane_x = []
|
||
plane_y = []
|
||
plane_z = []
|
||
print(box_list)
|
||
plane_x.append(box_list[0][0][0])
|
||
plane_x.append(box_list[0][1][0])
|
||
plane_x.append(box_list[0][2][0])
|
||
plane_x.append(box_list[0][3][0])
|
||
plane_y.append(box_list[0][0][1])
|
||
plane_y.append(box_list[0][1][1])
|
||
plane_y.append(box_list[0][2][1])
|
||
plane_y.append(box_list[0][3][1])
|
||
plane_z.append(box_list[0][0][2])
|
||
plane_z.append(box_list[0][1][2])
|
||
plane_z.append(box_list[0][2][2])
|
||
plane_z.append(box_list[0][3][2])
|
||
# 定义平面方程的参数
|
||
a = normal_vector[0]
|
||
b = normal_vector[1]
|
||
c = normal_vector[2]
|
||
d = normal_vector[3]
|
||
|
||
# 定义平面的范围
|
||
x_range = (int(min(plane_x)), int(max(plane_x)))
|
||
y_range = (int(min(plane_y)), int(max(plane_y)))
|
||
z_range = (int(min(plane_z)), int(max(plane_z)))
|
||
|
||
# 生成平面网格
|
||
x = np.linspace(x_range[0], x_range[1], 10)
|
||
y = np.linspace(y_range[0], y_range[1], 20)
|
||
X, Y = np.meshgrid(x, y)
|
||
Z = -(a * X + b * Y + d) / c # 根据平面方程计算 Z 坐标
|
||
|
||
# 确保 Z 坐标在指定范围内
|
||
Z = np.clip(Z, z_range[0], z_range[1])
|
||
|
||
# 创建 TriangleMesh 对象
|
||
import open3d as o3d
|
||
plane_mesh = o3d.geometry.TriangleMesh()
|
||
plane_mesh.vertices = o3d.utility.Vector3dVector(np.vstack((X.ravel(), Y.ravel(), Z.ravel())).T)
|
||
plane_mesh.triangles = o3d.utility.Vector3iVector(
|
||
np.array([[i, i + 1, i + 100] for i in range(99)] + [[i + 1, i + 101, i + 100] for i in range(99)]))
|
||
plane_mesh.paint_uniform_color([1, 0.5, 0.5]) # 设置平面颜色
|
||
|
||
return plane_mesh
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|