286 lines
8.1 KiB
Python
286 lines
8.1 KiB
Python
|
import json
|
||
|
|
import time
|
||
|
|
|
||
|
|
import cv2
|
||
|
|
import open3d as o3d
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
import transforms3d as tfs
|
||
|
|
|
||
|
|
from ultralytics import YOLO
|
||
|
|
|
||
|
|
import rclpy
|
||
|
|
from rclpy.node import Node
|
||
|
|
|
||
|
|
|
||
|
|
# x=0.08685893262849026, y=0.03733333467952511, z=0.5707736129272428
|
||
|
|
# x=-0.04385456738166652, y=-0.6302417335103124, z=-0.07781808527160532, w=0.7712434634188767
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
color_image = cv2.imread("test/color_image.png")
|
||
|
|
depth_image = cv2.imread("test/depth_image.png", cv2.IMREAD_UNCHANGED)
|
||
|
|
with open("test/K.json", 'r') as f:
|
||
|
|
k = json.load(f)
|
||
|
|
with open("test/D.json", 'r') as f:
|
||
|
|
d = json.load(f)
|
||
|
|
with open("test/eye_in_right_hand.json", 'r') as f:
|
||
|
|
hand_eye_mat = np.array(json.load(f)["T"])
|
||
|
|
|
||
|
|
camera_size = [color_image.shape[1], color_image.shape[0]]
|
||
|
|
|
||
|
|
# cv2.imshow("color_image", color_image)
|
||
|
|
# cv2.waitKey(0)
|
||
|
|
# cv2.destroyAllWindows()
|
||
|
|
# cv2.imshow("depth_image", depth_image)
|
||
|
|
# cv2.waitKey(0)
|
||
|
|
# cv2.destroyAllWindows()
|
||
|
|
# print(k)
|
||
|
|
# print(d)
|
||
|
|
# print(eye_in_right_hand)
|
||
|
|
|
||
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
|
|
model = YOLO("test/yolo11s-seg.pt").to(device)
|
||
|
|
|
||
|
|
results = model(color_image, retina_masks=True, conf=0.5, classes=[39])
|
||
|
|
result = results[0]
|
||
|
|
|
||
|
|
# Get masks
|
||
|
|
if result.masks is None or len(result.masks) == 0:
|
||
|
|
return None, None
|
||
|
|
masks = result.masks.data.cpu().numpy()
|
||
|
|
|
||
|
|
# Get boxes
|
||
|
|
boxes = result.boxes.xywh.cpu().numpy()
|
||
|
|
class_ids = result.boxes.cls.cpu().numpy()
|
||
|
|
labels = result.names
|
||
|
|
print(f"Detect object num: {len(masks)}")
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
for i, (mask, box) in enumerate(zip(masks, boxes)):
|
||
|
|
imgs_crop, (x_min, y_min) = crop_imgs_box_xywh([depth_image, mask], box, True)
|
||
|
|
depth_crop, mask_crop = imgs_crop
|
||
|
|
|
||
|
|
if depth_crop is None:
|
||
|
|
continue
|
||
|
|
|
||
|
|
intrinsics = o3d.camera.PinholeCameraIntrinsic(
|
||
|
|
int(camera_size[0]),
|
||
|
|
int(camera_size[1]),
|
||
|
|
k[0],
|
||
|
|
k[4],
|
||
|
|
k[2] - x_min,
|
||
|
|
k[5] - y_min
|
||
|
|
)
|
||
|
|
time1 = time.time()
|
||
|
|
rmat = calculate_pose_pca(
|
||
|
|
mask_crop,
|
||
|
|
depth_crop,
|
||
|
|
intrinsics,
|
||
|
|
)
|
||
|
|
time2 = time.time()
|
||
|
|
print(f"PCA: {(time2 - time1) * 1000} ms")
|
||
|
|
|
||
|
|
grab_width = calculate_grav_width(mask.astype(np.uint8) * 255, k, rmat[2, 3])
|
||
|
|
rmat[2, 3] = rmat[2, 3] + grab_width * 0.22
|
||
|
|
|
||
|
|
rmat = hand_eye_mat @ rmat
|
||
|
|
|
||
|
|
x, y, z, rw, rx, ry, rz = rmat2quat(rmat)
|
||
|
|
|
||
|
|
pose_list = []
|
||
|
|
if (x, y, z) != (0.0, 0.0, 0.0):
|
||
|
|
pose = [x, y, z, rx, ry, rz, rw]
|
||
|
|
pose_list.append(
|
||
|
|
{
|
||
|
|
"class_id": int(class_ids[i]),
|
||
|
|
"class_name": labels[class_ids[i]],
|
||
|
|
"pose": pose,
|
||
|
|
"grab_width": grab_width * 1.05
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def crop_imgs_box_xywh(imgs: list, box, same_sign: bool = False):
|
||
|
|
"""
|
||
|
|
Crop imgs
|
||
|
|
|
||
|
|
input:
|
||
|
|
imgs: list, Each img in imgs has the same Width and High.
|
||
|
|
box: The YOLO model outputs bounding box data in the format [x, y, w, h, confidence, class_id].
|
||
|
|
same_sign: bool, Set True to skip size check if all img in imgs have the same Width and High.
|
||
|
|
output:
|
||
|
|
crop_imgs: list;
|
||
|
|
(x_min, y_min);
|
||
|
|
"""
|
||
|
|
if not imgs:
|
||
|
|
print("imgs is empty")
|
||
|
|
return [], (0, 0)
|
||
|
|
|
||
|
|
if not same_sign and len(imgs) != 1:
|
||
|
|
for img in imgs:
|
||
|
|
if imgs[0].shape != img.shape:
|
||
|
|
raise ValueError(f"Img shape are different: {imgs[0].shape} - {img.shape}")
|
||
|
|
|
||
|
|
high, width = imgs[0].shape[:2]
|
||
|
|
x_center, y_center, w, h = box[:4]
|
||
|
|
x_min, x_max = max(0, int(round(x_center - w/2))), min(int(round(x_center + w/2)), width-1)
|
||
|
|
y_min, y_max = max(0, int(round(y_center - h/2))), min(int(round(y_center + h/2)), high-1)
|
||
|
|
|
||
|
|
crop_imgs = [img[y_min:y_max + 1, x_min:x_max + 1] for img in imgs]
|
||
|
|
|
||
|
|
return crop_imgs, (x_min, y_min)
|
||
|
|
|
||
|
|
|
||
|
|
def pca(data: np.ndarray, sort=True):
|
||
|
|
"""主成分分析 """
|
||
|
|
center = np.mean(data, axis=0)
|
||
|
|
centered_points = data - center # 去中心化
|
||
|
|
|
||
|
|
try:
|
||
|
|
cov_matrix = np.cov(centered_points.T) # 转置
|
||
|
|
eigenvalues, eigenvectors = np.linalg.eigh(cov_matrix)
|
||
|
|
|
||
|
|
except np.linalg.LinAlgError:
|
||
|
|
return None, None
|
||
|
|
|
||
|
|
if sort:
|
||
|
|
sort = eigenvalues.argsort()[::-1] # 降序排列
|
||
|
|
eigenvalues = eigenvalues[sort] # 索引
|
||
|
|
eigenvectors = eigenvectors[:, sort]
|
||
|
|
|
||
|
|
return eigenvalues, eigenvectors
|
||
|
|
|
||
|
|
|
||
|
|
def calculate_pose_pca(
|
||
|
|
mask,
|
||
|
|
depth_img: np.ndarray,
|
||
|
|
intrinsics,
|
||
|
|
source=None,
|
||
|
|
configs=None,
|
||
|
|
**kwargs
|
||
|
|
):
|
||
|
|
"""计算位态"""
|
||
|
|
if configs is None:
|
||
|
|
configs = {
|
||
|
|
"depth_scale": 1000.0,
|
||
|
|
"depth_trunc": 3.0,
|
||
|
|
"voxel_size": 0.010,
|
||
|
|
"nb_points": 16,
|
||
|
|
"radius": 0.03,
|
||
|
|
"nb_neighbors": 10,
|
||
|
|
"std_ratio": 3.0
|
||
|
|
}
|
||
|
|
depth_img_mask = np.zeros_like(depth_img)
|
||
|
|
depth_img_mask[mask > 0] = depth_img[mask > 0]
|
||
|
|
|
||
|
|
depth_o3d = o3d.geometry.Image(depth_img_mask.astype(np.uint16))
|
||
|
|
|
||
|
|
point_cloud = o3d.geometry.PointCloud.create_from_depth_image(
|
||
|
|
depth=depth_o3d,
|
||
|
|
intrinsic=intrinsics,
|
||
|
|
depth_scale=configs["depth_scale"],
|
||
|
|
depth_trunc=configs["depth_trunc"],
|
||
|
|
)
|
||
|
|
|
||
|
|
point_cloud = point_cloud.remove_non_finite_points()
|
||
|
|
|
||
|
|
down_pcd = point_cloud.voxel_down_sample(voxel_size=configs["voxel_size"])
|
||
|
|
clean_pcd, _ = down_pcd.remove_statistical_outlier(
|
||
|
|
nb_neighbors=configs["nb_neighbors"],
|
||
|
|
std_ratio=configs["std_ratio"]
|
||
|
|
)
|
||
|
|
clean_pcd, _ = clean_pcd.remove_radius_outlier(
|
||
|
|
nb_points=configs["nb_points"],
|
||
|
|
radius=configs["radius"]
|
||
|
|
)
|
||
|
|
labels = np.array(clean_pcd.cluster_dbscan(eps=0.03, min_points=16))
|
||
|
|
largest_cluster_idx = np.argmax(np.bincount(labels[labels >= 0]))
|
||
|
|
clean_pcd = clean_pcd.select_by_index(np.where(labels == largest_cluster_idx)[0])
|
||
|
|
|
||
|
|
|
||
|
|
if len(clean_pcd.points) == 0:
|
||
|
|
print("clean_pcd is empty")
|
||
|
|
return 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
|
||
|
|
center = clean_pcd.get_center()
|
||
|
|
x, y, z = center
|
||
|
|
|
||
|
|
w, v = pca(np.asarray(clean_pcd.points))
|
||
|
|
|
||
|
|
if w is None or v is None:
|
||
|
|
print("PCA output w or v is None")
|
||
|
|
return np.eye(4)
|
||
|
|
|
||
|
|
vx, vy, vz = v[:, 0], v[:, 1], v[:, 2]
|
||
|
|
|
||
|
|
if vx[0] < 0:
|
||
|
|
vx = -vx
|
||
|
|
if vy[1] < 0:
|
||
|
|
vy = -vy
|
||
|
|
if not np.allclose(np.cross(vx, vy), vz):
|
||
|
|
vz = -vz
|
||
|
|
|
||
|
|
R = np.column_stack((vx, vy, vz))
|
||
|
|
rmat = tfs.affines.compose(np.squeeze(np.asarray((x, y, z+0.06*0.25))), R, [1, 1, 1])
|
||
|
|
|
||
|
|
draw(point_cloud, rmat)
|
||
|
|
# draw(down_pcd, rmat)
|
||
|
|
draw(clean_pcd, rmat)
|
||
|
|
|
||
|
|
return rmat
|
||
|
|
|
||
|
|
|
||
|
|
def rmat2quat(rmat):
|
||
|
|
"""Convert rotation matrix to quaternion."""
|
||
|
|
x, y, z = rmat[0:3, 3:4].flatten()
|
||
|
|
rw, rx, ry, rz = tfs.quaternions.mat2quat(rmat[0:3, 0:3])
|
||
|
|
quat = [x, y, z, rw, rx, ry, rz]
|
||
|
|
return quat
|
||
|
|
|
||
|
|
|
||
|
|
def calculate_grav_width(mask, k, depth):
|
||
|
|
"""计算重心宽度"""
|
||
|
|
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
|
|
|
||
|
|
if not contours:
|
||
|
|
return 0.0
|
||
|
|
|
||
|
|
c = max(contours, key=cv2.contourArea)
|
||
|
|
box = cv2.boxPoints(cv2.minAreaRect(c))
|
||
|
|
if np.linalg.norm(box[1] - box[0]) < np.linalg.norm(box[1] - box[2]):
|
||
|
|
point_diff = box[1] - box[0]
|
||
|
|
else:
|
||
|
|
point_diff = box[1] - box[2]
|
||
|
|
|
||
|
|
grab_width = depth * np.sqrt(
|
||
|
|
point_diff[0] ** 2 / k[0] ** 2 + point_diff[1] ** 2 / k[4] ** 2
|
||
|
|
)
|
||
|
|
return grab_width
|
||
|
|
|
||
|
|
def draw(pcd, mat):
|
||
|
|
R = mat[0:3, 0:3]
|
||
|
|
point = mat[0:3, 3:4].flatten()
|
||
|
|
x, y, z = R[:, 0], R[:, 1], R[:, 2]
|
||
|
|
|
||
|
|
points = [
|
||
|
|
[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1],
|
||
|
|
point, point + x, point + y, point + z
|
||
|
|
|
||
|
|
] # 画点：原点、第一主成分、第二主成分
|
||
|
|
lines = [
|
||
|
|
[0, 1], [0, 2], [0, 3],
|
||
|
|
[4, 5], [4, 6], [4, 7]
|
||
|
|
] # 画出三点之间两两连线
|
||
|
|
colors = [
|
||
|
|
[1, 0, 0], [0, 1, 0], [0, 0, 1],
|
||
|
|
[1, 0, 0], [0, 1, 0], [0, 0, 1]
|
||
|
|
]
|
||
|
|
line_set = o3d.geometry.LineSet(points=o3d.utility.Vector3dVector(points), lines=o3d.utility.Vector2iVector(lines))
|
||
|
|
line_set.colors = o3d.utility.Vector3dVector(colors)
|
||
|
|
|
||
|
|
o3d.visualization.draw_geometries([pcd, line_set])
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
main()
|