6 | [](https://arxiv.org/abs/2402.05655.pdf)
7 | 
8 | 
9 |
10 | [](https://paperswithcode.com/sota/robot-pose-estimation-on-dream-dataset?p=real-time-holistic-robot-pose-estimation-with)
11 |
12 | 
15 |
16 | This is the official PyTorch implementation of the paper "Real-time Holistic Robot Pose Estimation with Unknown States". It provides an efficient framework for real-time robot pose estimation from RGB images without requiring known robot states.
17 |
18 | ## Installation
19 | This project's dependencies include python 3.9, pytorch 1.13, pytorch3d 0.7.4 and CUDA 11.7.
20 | The code is developed and tested on Ubuntu 20.04.
21 |
22 | ```bash
23 | pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117
24 | pip install -r requirements.txt
25 | conda install pytorch3d=0.7.4 # from https://anaconda.org/pytorch3d/pytorch3d/files
26 | ```
27 |
28 | ## Data and Model Preparation
29 |
30 | In our work, we use the following data and pretrained model:
31 | * The [DREAM datasets](https://drive.google.com/drive/folders/1uNK2n9wU4tRE07sM_r640wDhwmOwuxx6) consisting of both real and synthetic subsets, placed under `${ROOT}/data/dream/$`.
32 | * The [URDF](https://drive.google.com/drive/folders/17KNhy28pypheYfDCxgOjJf4IyUnOI3gW?) (Unified Robotics Description Format) of robot Panda, Kuka and Baxter, placed under `${ROOT}/data/deps/$`.
33 | * The [pretrained HRnet backbone](https://drive.google.com/file/d/1eqIftq1T_oIGhmCfkVYSM245Wj5xZaUo/view?) for pose estimation, placed under `${ROOT}/models/$`.
34 | * The openly available [foreground segmentation model](https://drive.google.com/drive/folders/1PpXe3p5dJt9EOM-fwvJ9TNStTWTQFDNK?) of 4 real datasets of Panda from [CtRNet](https://github.com/ucsdarclab/CtRNet-robot-pose-estimation), placed under `${ROOT}/models/panda_segmentation/$`.
35 |
36 | You can download the data and models through provided links.
37 | When finished, the directory tree should look like this.
38 | ```
39 | ${ROOT}
40 | |-- data
41 | |-- dream
42 | | |-- real
43 | | | |-- panda-3cam_azure
44 | | | |-- panda-3cam_kinect360
45 | | | |-- panda-3cam_realsense
46 | | | |-- panda-orb
47 | | |-- synthetic
48 | | | |-- baxter_synth_test_dr
49 | | | |-- baxter_synth_train_dr
50 | | | |-- kuka_synth_test_dr
51 | | | |-- kuka_synth_test_photo
52 | | | |-- kuka_synth_train_dr
53 | | | |-- panda_synth_test_dr
54 | | | |-- panda_synth_test_photo
55 | | | |-- panda_synth_train_dr
56 | |-- deps
57 | | |-- baxter-description
58 | | |-- kuka-description
59 | | |-- panda-description
60 | |-- models
61 | |-- panda_segmentation
62 | | |-- azure.pth
63 | | |-- kinect.pth
64 | | |-- orb.pth
65 | | |-- realsense.pth
66 | |-- hrnet_w32-36af842e_roc.pth
67 | ```
68 |
69 | ## Train
70 | We train our final model in a multi-stage fashion. All model is trained using a single NVIDIA V100 with 32GB GPU. Distributed training is also supported.
71 |
72 | We use config files in `configs/` to specify the training process. We recommend filling in the `exp_name` field in the config files with a unique name, as the checkpoints and event logs produced during training will be saved under `experiments/{exp_name}`. The correspondent config file will be automatically copied into this directory.
73 |
74 | ### Synthetic Datasets
75 |
76 | Firstly, pretrain the depthnet (root depth estimator) for 100 epochs for each robot arm:
77 | ```bash
78 | python scripts/train.py --config configs/panda/depthnet.yaml
79 | python scripts/train.py --config configs/kuka/depthnet.yaml
80 | python scripts/train.py --config configs/baxter/depthnet.yaml
81 | ```
82 |
83 | With depthnet pretrained, we can train the full network for 100 epochs:
84 | ```bash
85 | python scripts/train.py --config configs/panda/full.yaml
86 | python scripts/train.py --config configs/kuka/full.yaml
87 | python scripts/train.py --config configs/baxter/full.yaml
88 | ```
89 | To save your time when reproducing results of our paper, we provide readily-pretrained [depthnet model weights](https://drive.google.com/drive/folders/1rWC2bbA3U0IiZ7oDoKIVsWK_m4JkVarA?) for full network training. To use them, you can modify the `configs/{robot}/full.yaml` file by filling in the `pretrained_rootnet` field with the path of the downloaded `.pk` file.
90 |
91 | ### Real Datasets of Panda
92 |
93 | We employ self-supervised training for the 4 real datasets of Panda.
94 |
95 | Firstly, train the model on synthetic dataset using `configs/panda/self_supervised/synth.yaml` for 100 epochs. Be sure to fill in the `pretrained_rootnet` field with the path of the pretrained Panda depthnet weight in advance.
96 |
97 | ```bash
98 | python scripts/train.py --config configs/panda/self_supervised/synth.yaml
99 | ```
100 | The training process above saves checkpoints for 4 real datasets for further self-supervised training (e.g. `experiments/{exp_name}/ckpt/curr_best_auc(add)_azure_model.pk`).
101 |
102 | When finished training on synthetic data, modify the `configs/panda/self_supervised/{real_dataset}.yaml` file by filling in the `pretrained_weight_on_synth` field with the path of the correspondent checkpoint. Then start self-supervised training with:
103 |
104 | ```bash
105 | python scripts/train.py --config configs/panda/self_supervised/azure.yaml
106 | python scripts/train.py --config configs/panda/self_supervised/kinect.yaml
107 | python scripts/train.py --config configs/panda/self_supervised/realsense.yaml
108 | python scripts/train.py --config configs/panda/self_supervised/orb.yaml
109 | ```
110 |
111 | ## Test
112 | To evaluate models, simply run:
113 | ```bash
114 | python scripts/test.py --exp_path {path of the experiment folder} --dataset {dataset name}
115 | # e.g. python scripts/test.py -e experiments/panda_full --dataset panda_synth_test_dr
116 | # You can add '--vis_skeleton' to visualize the robot keypoint skeleton
117 | ```
118 | Note that each model is presented in a folder containing ckpt/, log/ and config.yaml. After running test script, result/ will be generated inside the folder.
119 |
120 | ## Model Zoo
121 | You can download our final models from [Google Drive](https://drive.google.com/drive/folders/10Gz0NP39YyuvAlrhTa-XssWTDlyh9v80?usp=sharing) and evaluate them yourself.
122 |
123 |
124 | ## Citation
125 | If you use our code or models in your research, please cite with:
126 | ```bibtex
127 | @inproceedings{holisticrobotpose,
128 | author={Ban, Shikun and Fan, Juling and Ma, Xiaoxuan and Zhu, Wentao and Qiao, Yu and Wang, Yizhou},
129 | title={Real-time Holistic Robot Pose Estimation with Unknown States},
130 | booktitle = {European Conference on Computer Vision (ECCV)},
131 | year = {2024}
132 | }
133 | ```
134 |
135 | ## Acknowledgment
136 | This repo is built on the excellent work [RoboPose](https://github.com/ylabbe/robopose) and [CtRNet](https://github.com/ucsdarclab/CtRNet-robot-pose-estimation). Thank the authors for releasing their codes.
137 |
--------------------------------------------------------------------------------
/lib/utils/mesh_renderer.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import numpy as np
3 |
4 | # io utils
5 | from pytorch3d.io import load_obj
6 |
7 | # datastructures
8 | from pytorch3d.structures import Meshes
9 |
10 | # rendering components
11 | from pytorch3d.renderer import (
12 | RasterizationSettings, MeshRenderer, MeshRasterizer, BlendParams,
13 | SoftSilhouetteShader, HardPhongShader, PointLights, TexturesVertex,
14 | PerspectiveCameras,Textures
15 | )
16 |
17 | from os.path import exists
18 | from roboticstoolbox.robot.ERobot import ERobot
19 |
20 |
21 | class PandaArm():
22 | def __init__(self, urdf_file):
23 |
24 | self.robot = self.Panda(urdf_file)
25 |
26 | def get_joint_RT(self, joint_angle):
27 |
28 | assert joint_angle.shape[0] == 7
29 |
30 | link_idx_list = [0,1,2,3,4,5,6,7,9]
31 | # link 0,1,2,3,4,5,6,7, and hand
32 | R_list = []
33 | t_list = []
34 |
35 | for i in range(len(link_idx_list)):
36 | link_idx = link_idx_list[i]
37 | T = self.robot.fkine(joint_angle, end = self.robot.links[link_idx], start = self.robot.links[0])
38 | R_list.append(T.R)
39 | t_list.append(T.t)
40 |
41 | return np.array(R_list),np.array(t_list)
42 |
43 | class Panda(ERobot):
44 | """
45 | Class that imports a URDF model
46 | """
47 |
48 | def __init__(self, urdf_file):
49 |
50 | links, name, urdf_string, urdf_filepath = self.URDF_read(urdf_file)
51 |
52 | super().__init__(
53 | links,
54 | name=name,
55 | manufacturer="Franka",
56 | urdf_string=urdf_string,
57 | urdf_filepath=urdf_filepath,
58 | )
59 |
60 |
61 | class RobotMeshRenderer():
62 | """
63 | Class that render robot mesh with differentiable renderer
64 | """
65 | def __init__(self, focal_length, principal_point, image_size, robot, mesh_files, device):
66 |
67 | self.focal_length = focal_length
68 | self.principal_point = principal_point
69 | self.image_size = image_size
70 | self.device = device
71 | self.robot = robot
72 | self.mesh_files = mesh_files
73 | self.preload_verts = []
74 | self.preload_faces = []
75 |
76 |
77 | # preload the mesh to save loading time
78 | for m_file in mesh_files:
79 | assert exists(m_file)
80 | preload_verts_i, preload_faces_idx_i, _ = load_obj(m_file)
81 | preload_faces_i = preload_faces_idx_i.verts_idx
82 | self.preload_verts.append(preload_verts_i)
83 | self.preload_faces.append(preload_faces_i)
84 |
85 |
86 | # set up differentiable renderer with given camera parameters
87 | self.cameras = PerspectiveCameras(
88 | focal_length = [focal_length],
89 | principal_point = [principal_point],
90 | device=device,
91 | in_ndc=False, image_size = [image_size]
92 | ) # (height, width) !!!!!
93 |
94 | blend_params = BlendParams(sigma=1e-8, gamma=1e-8)
95 | raster_settings = RasterizationSettings(
96 | image_size=image_size,
97 | blur_radius=np.log(1. / 1e-4 - 1.) * blend_params.sigma,
98 | faces_per_pixel=100,
99 | max_faces_per_bin=100000, # max_faces_per_bin=1000000,
100 | )
101 |
102 | # Create a silhouette mesh renderer by composing a rasterizer and a shader.
103 | self.silhouette_renderer = MeshRenderer(
104 | rasterizer=MeshRasterizer(
105 | cameras=self.cameras,
106 | raster_settings=raster_settings
107 | ),
108 | shader=SoftSilhouetteShader(blend_params=blend_params)
109 | )
110 |
111 |
112 | # We will also create a Phong renderer. This is simpler and only needs to render one face per pixel.
113 | raster_settings = RasterizationSettings(
114 | image_size=image_size,
115 | blur_radius=0.0,
116 | faces_per_pixel=1,
117 | max_faces_per_bin=100000,
118 | )
119 | # We can add a point light in front of the object.
120 | lights = PointLights(device=device, location=((2.0, 2.0, -2.0),))
121 | self.phong_renderer = MeshRenderer(
122 | rasterizer=MeshRasterizer(
123 | cameras=self.cameras,
124 | raster_settings=raster_settings
125 | ),
126 | shader=HardPhongShader(device=device, cameras=self.cameras, lights=lights)
127 | )
128 |
129 | def get_robot_mesh(self, joint_angle):
130 |
131 | R_list, t_list = self.robot.get_joint_RT(joint_angle)
132 | assert len(self.mesh_files) == R_list.shape[0] and len(self.mesh_files) == t_list.shape[0]
133 |
134 | verts_list = []
135 | faces_list = []
136 | verts_rgb_list = []
137 | verts_count = 0
138 | for i in range(len(self.mesh_files)):
139 | verts_i = self.preload_verts[i]
140 | faces_i = self.preload_faces[i]
141 |
142 | R = torch.tensor(R_list[i],dtype=torch.float32)
143 | t = torch.tensor(t_list[i],dtype=torch.float32)
144 | verts_i = verts_i @ R.T + t
145 | #verts_i = (R @ verts_i.T).T + t
146 | faces_i = faces_i + verts_count
147 |
148 | verts_count+=verts_i.shape[0]
149 |
150 | verts_list.append(verts_i.to(self.device))
151 | faces_list.append(faces_i.to(self.device))
152 |
153 | # Initialize each vertex to be white in color.
154 | color = torch.rand(3)
155 | verts_rgb_i = torch.ones_like(verts_i) * color # (V, 3)
156 | verts_rgb_list.append(verts_rgb_i.to(self.device))
157 |
158 |
159 |
160 | verts = torch.concat(verts_list, dim=0)
161 | faces = torch.concat(faces_list, dim=0)
162 |
163 | verts_rgb = torch.concat(verts_rgb_list,dim=0)[None]
164 | textures = Textures(verts_rgb=verts_rgb)
165 |
166 | # Create a Meshes object
167 | robot_mesh = Meshes(
168 | verts=[verts.to(self.device)],
169 | faces=[faces.to(self.device)],
170 | textures=textures
171 | )
172 |
173 | return robot_mesh
174 |
175 |
176 | def get_robot_verts_and_faces(self, joint_angle):
177 |
178 | R_list, t_list = self.robot.get_joint_RT(joint_angle)
179 | assert len(self.mesh_files) == R_list.shape[0] and len(self.mesh_files) == t_list.shape[0]
180 |
181 | verts_list = []
182 | faces_list = []
183 | verts_rgb_list = []
184 | verts_count = 0
185 | for i in range(len(self.mesh_files)):
186 | verts_i = self.preload_verts[i]
187 | faces_i = self.preload_faces[i]
188 |
189 | R = torch.tensor(R_list[i],dtype=torch.float32)
190 | t = torch.tensor(t_list[i],dtype=torch.float32)
191 | verts_i = verts_i @ R.T + t
192 | #verts_i = (R @ verts_i.T).T + t
193 | faces_i = faces_i + verts_count
194 |
195 | verts_count+=verts_i.shape[0]
196 |
197 | verts_list.append(verts_i.to(self.device))
198 | faces_list.append(faces_i.to(self.device))
199 |
200 | # Initialize each vertex to be white in color.
201 | #color = torch.rand(3)
202 | #verts_rgb_i = torch.ones_like(verts_i) * color # (V, 3)
203 | #verts_rgb_list.append(verts_rgb_i.to(self.device))
204 |
205 | verts = torch.concat(verts_list, dim=0)
206 | faces = torch.concat(faces_list, dim=0)
207 |
208 |
209 | return verts, faces
--------------------------------------------------------------------------------
/lib/models/ctrnet/CtRNet.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn.functional as F
3 | import kornia
4 | import numpy as np
5 |
6 | from .keypoint_seg_resnet import KeyPointSegNet
7 | from utils.BPnP import BPnP, BPnP_m3d, batch_project
8 |
9 |
10 | class CtRNet(torch.nn.Module):
11 | def __init__(self, args):
12 | super(CtRNet, self).__init__()
13 |
14 | self.args = args
15 |
16 | if args.use_gpu:
17 | self.device = "cuda"
18 | else:
19 | self.device = "cpu"
20 |
21 | # load keypoint segmentation model
22 | self.keypoint_seg_predictor = KeyPointSegNet(args, use_gpu=args.use_gpu)
23 |
24 | if args.use_gpu:
25 | self.keypoint_seg_predictor = self.keypoint_seg_predictor.cuda()
26 |
27 | self.keypoint_seg_predictor = torch.nn.DataParallel(self.keypoint_seg_predictor, device_ids=[0])
28 |
29 | if args.keypoint_seg_model_path is not None:
30 | print("Loading keypoint segmentation model from {}".format(args.keypoint_seg_model_path))
31 | self.keypoint_seg_predictor.load_state_dict(torch.load(args.keypoint_seg_model_path))
32 |
33 | self.keypoint_seg_predictor.eval()
34 |
35 | # load BPnP
36 | self.bpnp = BPnP.apply
37 | self.bpnp_m3d = BPnP_m3d.apply
38 |
39 | # set up camera intrinsics
40 |
41 | self.intrinsics = np.array([[ args.fx, 0. , args.px ],
42 | [ 0. , args.fy, args.py ],
43 | [ 0. , 0. , 1. ]])
44 | print("Camera intrinsics: {}".format(self.intrinsics))
45 |
46 | self.K = torch.tensor(self.intrinsics, device=self.device, dtype=torch.float)
47 |
48 |
49 | def inference_single_image(self, img, joint_angles):
50 | # img: (3, H, W)
51 | # joint_angles: (7)
52 | # robot: robot model
53 |
54 | # detect 2d keypoints and segmentation masks
55 | points_2d, segmentation = self.keypoint_seg_predictor(img[None])
56 | foreground_mask = torch.sigmoid(segmentation)
57 | _,t_list = self.robot.get_joint_RT(joint_angles)
58 | points_3d = torch.from_numpy(np.array(t_list)).float().to(self.device)
59 | if self.args.robot_name == "Panda":
60 | points_3d = points_3d[[0,2,3,4,6,7,8]] # remove 1 and 5 links as they are overlapping with 2 and 6
61 |
62 | #init_pose = torch.tensor([[ 1.5497, 0.5420, -0.3909, -0.4698, -0.0211, 1.3243]])
63 | #cTr = bpnp(points_2d_pred, points_3d, K, init_pose)
64 | cTr = self.bpnp(points_2d, points_3d, self.K)
65 |
66 | return cTr, points_2d, foreground_mask
67 |
68 | def inference_batch_images(self, img, joint_angles):
69 | # img: (B, 3, H, W)
70 | # joint_angles: (B, 7)
71 | # robot: robot model
72 |
73 | # detect 2d keypoints and segmentation masks
74 | points_2d, segmentation = self.keypoint_seg_predictor(img)
75 | foreground_mask = torch.sigmoid(segmentation)
76 |
77 | points_3d_batch = []
78 | for b in range(joint_angles.shape[0]):
79 | _,t_list = self.robot.get_joint_RT(joint_angles[b])
80 | points_3d = torch.from_numpy(np.array(t_list)).float().to(self.device)
81 | if self.args.robot_name == "Panda":
82 | points_3d = points_3d[:,[0,2,3,4,6,7,8]]
83 | points_3d_batch.append(points_3d[None])
84 |
85 | points_3d_batch = torch.cat(points_3d_batch, dim=0)
86 |
87 | cTr = self.bpnp_m3d(points_2d, points_3d_batch, self.K)
88 |
89 | return cTr, points_2d, foreground_mask
90 |
91 | def inference_batch_images_seg_kp(self, img):
92 | # img: (B, 3, H, W)
93 | # joint_angles: (B, 7)
94 | # robot: robot model
95 |
96 | # detect 2d keypoints and segmentation masks
97 | points_2d, segmentation = self.keypoint_seg_predictor(img)
98 | foreground_mask = torch.sigmoid(segmentation)
99 |
100 | return points_2d, foreground_mask
101 |
102 | def inference_batch_images_onlyseg(self, img):
103 | # img: (B, 3, H, W)
104 | # joint_angles: (B, 7)
105 | # robot: robot model
106 |
107 | # detect 2d keypoints and segmentation masks
108 | points_2d, segmentation = self.keypoint_seg_predictor(img)
109 | foreground_mask = torch.sigmoid(segmentation)
110 |
111 | return foreground_mask
112 |
113 |
114 | def cTr_to_pose_matrix(self, cTr):
115 | """
116 | cTr: (batch_size, 6)
117 | pose_matrix: (batch_size, 4, 4)
118 | """
119 | batch_size = cTr.shape[0]
120 | pose_matrix = torch.zeros((batch_size, 4, 4), device=self.device)
121 | pose_matrix[:, :3, :3] = kornia.geometry.conversions.axis_angle_to_rotation_matrix(cTr[:, :3])
122 | pose_matrix[:, :3, 3] = cTr[:, 3:]
123 | pose_matrix[:, 3, 3] = 1
124 | return pose_matrix
125 |
126 | def to_valid_R_batch(self, R):
127 | # R is a batch of 3x3 rotation matrices
128 | U, S, V = torch.svd(R)
129 | return torch.bmm(U, V.transpose(1,2))
130 |
131 | def render_single_robot_mask(self, cTr, robot_mesh, robot_renderer):
132 | # cTr: (6)
133 | # img: (1, H, W)
134 |
135 | R = kornia.geometry.conversions.angle_axis_to_rotation_matrix(cTr[:3][None]) # (1, 3, 3)
136 | R = torch.transpose(R,1,2)
137 | #R = to_valid_R_batch(R)
138 | T = cTr[3:][None] # (1, 3)
139 |
140 | if T[0,-1] < 0:
141 | rendered_image = robot_renderer.silhouette_renderer(meshes_world=robot_mesh, R = -R, T = -T)
142 | else:
143 | rendered_image = robot_renderer.silhouette_renderer(meshes_world=robot_mesh, R = R, T = T)
144 |
145 | if torch.isnan(rendered_image).any():
146 | rendered_image = torch.nan_to_num(rendered_image)
147 |
148 | return rendered_image[..., 3]
149 |
150 |
151 | def train_on_batch(self, img, joint_angles, robot_renderer, criterions, phase='train'):
152 | # img: (B, 3, H, W)
153 | # joint_angles: (B, 7)
154 | with torch.set_grad_enabled(phase == 'train'):
155 | # detect 2d keypoints
156 | points_2d, segmentation = self.keypoint_seg_predictor(img)
157 |
158 | mask_list = list()
159 | seg_weight_list = list()
160 |
161 | for b in range(img.shape[0]):
162 | # get 3d points
163 | _,t_list = self.robot.get_joint_RT(joint_angles[b])
164 | points_3d = torch.from_numpy(np.array(t_list)).float().to(self.device)
165 | if self.args.robot_name == "Panda":
166 | points_3d = points_3d[:,[0,2,3,4,6,7,8]]
167 |
168 | # get camera pose
169 | cTr = self.bpnp(points_2d[b][None], points_3d, self.K)
170 |
171 | # config robot mesh
172 | robot_mesh = robot_renderer.get_robot_mesh(joint_angles[b])
173 |
174 | # render robot mask
175 | rendered_image = self.render_single_robot_mask(cTr.squeeze(), robot_mesh, robot_renderer)
176 |
177 | mask_list.append(rendered_image)
178 | points_2d_proj = batch_project(cTr, points_3d, self.K)
179 | reproject_error = criterions["mse_mean"](points_2d[b], points_2d_proj.squeeze())
180 | seg_weight = torch.exp(-reproject_error * self.args.reproj_err_scale)
181 | seg_weight_list.append(seg_weight)
182 |
183 | mask_batch = torch.cat(mask_list,0)
184 |
185 | loss_bce = 0
186 | for b in range(segmentation.shape[0]):
187 | loss_bce = loss_bce + seg_weight_list[b] * criterions["bce"](segmentation[b].squeeze(), mask_batch[b].detach())
188 |
189 | img_ref = torch.sigmoid(segmentation).detach()
190 | #loss_reproj = 0.0005 * criterionMSE_mean(points_2d, points_2d_proj_batch)
191 | loss_mse = 0.001 * criterions["mse_sum"](mask_batch, img_ref.squeeze())
192 | loss = loss_mse + loss_bce
193 |
194 | return loss
195 |
196 |
197 |
198 |
199 |
200 |
--------------------------------------------------------------------------------
/lib/utils/urdfpytorch/utils.py:
--------------------------------------------------------------------------------
1 | """Utilities for URDF parsing.
2 | """
3 | import os
4 | from pathlib import Path
5 |
6 | from lxml import etree as ET
7 | import numpy as np
8 | import trimesh
9 |
10 |
11 | def resolve_package_path(urdf_path, mesh_path):
12 | urdf_path = Path(urdf_path)
13 | search_dir = urdf_path
14 | relative_path = Path(str(mesh_path).replace('package://', ''))
15 | while True:
16 | absolute_path = (search_dir / relative_path)
17 | if absolute_path.exists():
18 | return absolute_path
19 | search_dir = search_dir.parent
20 |
21 |
22 | def rpy_to_matrix(coords):
23 | """Convert roll-pitch-yaw coordinates to a 3x3 homogenous rotation matrix.
24 |
25 | The roll-pitch-yaw axes in a typical URDF are defined as a
26 | rotation of ``r`` radians around the x-axis followed by a rotation of
27 | ``p`` radians around the y-axis followed by a rotation of ``y`` radians
28 | around the z-axis. These are the Z1-Y2-X3 Tait-Bryan angles. See
29 | Wikipedia_ for more information.
30 |
31 | .. _Wikipedia: https://en.wikipedia.org/wiki/Euler_angles#Rotation_matrix
32 |
33 | Parameters
34 | ----------
35 | coords : (3,) float
36 | The roll-pitch-yaw coordinates in order (x-rot, y-rot, z-rot).
37 |
38 | Returns
39 | -------
40 | R : (3,3) float
41 | The corresponding homogenous 3x3 rotation matrix.
42 | """
43 | coords = np.asanyarray(coords, dtype=np.float64)
44 | c3, c2, c1 = np.cos(coords)
45 | s3, s2, s1 = np.sin(coords)
46 |
47 | return np.array([
48 | [c1 * c2, (c1 * s2 * s3) - (c3 * s1), (s1 * s3) + (c1 * c3 * s2)],
49 | [c2 * s1, (c1 * c3) + (s1 * s2 * s3), (c3 * s1 * s2) - (c1 * s3)],
50 | [-s2, c2 * s3, c2 * c3]
51 | ], dtype=np.float64)
52 |
53 |
54 | def matrix_to_rpy(R, solution=1):
55 | """Convert a 3x3 transform matrix to roll-pitch-yaw coordinates.
56 |
57 | The roll-pitchRyaw axes in a typical URDF are defined as a
58 | rotation of ``r`` radians around the x-axis followed by a rotation of
59 | ``p`` radians around the y-axis followed by a rotation of ``y`` radians
60 | around the z-axis. These are the Z1-Y2-X3 Tait-Bryan angles. See
61 | Wikipedia_ for more information.
62 |
63 | .. _Wikipedia: https://en.wikipedia.org/wiki/Euler_angles#Rotation_matrix
64 |
65 | There are typically two possible roll-pitch-yaw coordinates that could have
66 | created a given rotation matrix. Specify ``solution=1`` for the first one
67 | and ``solution=2`` for the second one.
68 |
69 | Parameters
70 | ----------
71 | R : (3,3) float
72 | A 3x3 homogenous rotation matrix.
73 | solution : int
74 | Either 1 or 2, indicating which solution to return.
75 |
76 | Returns
77 | -------
78 | coords : (3,) float
79 | The roll-pitch-yaw coordinates in order (x-rot, y-rot, z-rot).
80 | """
81 | R = np.asanyarray(R, dtype=np.float64)
82 | r = 0.0
83 | p = 0.0
84 | y = 0.0
85 |
86 | if np.abs(R[2,0]) >= 1.0 - 1e-12:
87 | y = 0.0
88 | if R[2,0] < 0:
89 | p = np.pi / 2
90 | r = np.arctan2(R[0,1], R[0,2])
91 | else:
92 | p = -np.pi / 2
93 | r = np.arctan2(-R[0,1], -R[0,2])
94 | else:
95 | if solution == 1:
96 | p = -np.arcsin(R[2,0])
97 | else:
98 | p = np.pi + np.arcsin(R[2,0])
99 | r = np.arctan2(R[2,1] / np.cos(p), R[2,2] / np.cos(p))
100 | y = np.arctan2(R[1,0] / np.cos(p), R[0,0] / np.cos(p))
101 |
102 | return np.array([r, p, y], dtype=np.float64)
103 |
104 |
105 | def matrix_to_xyz_rpy(matrix):
106 | """Convert a 4x4 homogenous matrix to xyzrpy coordinates.
107 |
108 | Parameters
109 | ----------
110 | matrix : (4,4) float
111 | The homogenous transform matrix.
112 |
113 | Returns
114 | -------
115 | xyz_rpy : (6,) float
116 | The xyz_rpy vector.
117 | """
118 | xyz = matrix[:3,3]
119 | rpy = matrix_to_rpy(matrix[:3,:3])
120 | return np.hstack((xyz, rpy))
121 |
122 |
123 | def xyz_rpy_to_matrix(xyz_rpy):
124 | """Convert xyz_rpy coordinates to a 4x4 homogenous matrix.
125 |
126 | Parameters
127 | ----------
128 | xyz_rpy : (6,) float
129 | The xyz_rpy vector.
130 |
131 | Returns
132 | -------
133 | matrix : (4,4) float
134 | The homogenous transform matrix.
135 | """
136 | matrix = np.eye(4, dtype=np.float64)
137 | matrix[:3,3] = xyz_rpy[:3]
138 | matrix[:3,:3] = rpy_to_matrix(xyz_rpy[3:])
139 | return matrix
140 |
141 |
142 | def parse_origin(node):
143 | """Find the ``origin`` subelement of an XML node and convert it
144 | into a 4x4 homogenous transformation matrix.
145 |
146 | Parameters
147 | ----------
148 | node : :class`lxml.etree.Element`
149 | An XML node which (optionally) has a child node with the ``origin``
150 | tag.
151 |
152 | Returns
153 | -------
154 | matrix : (4,4) float
155 | The 4x4 homogneous transform matrix that corresponds to this node's
156 | ``origin`` child. Defaults to the identity matrix if no ``origin``
157 | child was found.
158 | """
159 | matrix = np.eye(4, dtype=np.float64)
160 | origin_node = node.find('origin')
161 | if origin_node is not None:
162 | if 'xyz' in origin_node.attrib:
163 | matrix[:3,3] = np.fromstring(origin_node.attrib['xyz'], sep=' ')
164 | if 'rpy' in origin_node.attrib:
165 | rpy = np.fromstring(origin_node.attrib['rpy'], sep=' ')
166 | matrix[:3,:3] = rpy_to_matrix(rpy)
167 | return matrix
168 |
169 |
170 | def unparse_origin(matrix):
171 | """Turn a 4x4 homogenous matrix into an ``origin`` XML node.
172 |
173 | Parameters
174 | ----------
175 | matrix : (4,4) float
176 | The 4x4 homogneous transform matrix to convert into an ``origin``
177 | XML node.
178 |
179 | Returns
180 | -------
181 | node : :class`lxml.etree.Element`
182 | An XML node whose tag is ``origin``. The node has two attributes:
183 |
184 | - ``xyz`` - A string with three space-delimited floats representing
185 | the translation of the origin.
186 | - ``rpy`` - A string with three space-delimited floats representing
187 | the rotation of the origin.
188 | """
189 | node = ET.Element('origin')
190 | node.attrib['xyz'] = '{} {} {}'.format(*matrix[:3,3])
191 | node.attrib['rpy'] = '{} {} {}'.format(*matrix_to_rpy(matrix[:3,:3]))
192 | return node
193 |
194 |
195 | def get_filename(base_path, file_path, makedirs=False):
196 | """Formats a file path correctly for URDF loading.
197 |
198 | Parameters
199 | ----------
200 | base_path : str
201 | The base path to the URDF's folder.
202 | file_path : str
203 | The path to the file.
204 | makedirs : bool, optional
205 | If ``True``, the directories leading to the file will be created
206 | if needed.
207 |
208 | Returns
209 | -------
210 | resolved : str
211 | The resolved filepath -- just the normal ``file_path`` if it was an
212 | absolute path, otherwise that path joined to ``base_path``.
213 | """
214 | # print(base_path)
215 | # print(file_path)
216 | fn = file_path
217 | if not os.path.isabs(file_path):
218 | fn = os.path.join(base_path, file_path)
219 | if makedirs:
220 | d, _ = os.path.split(fn)
221 | if not os.path.exists(d):
222 | os.makedirs(d)
223 | if not Path(fn).exists():
224 | fn = str(resolve_package_path(base_path, file_path))
225 | return fn
226 |
227 |
228 | def load_meshes(filename):
229 | """Loads triangular meshes from a file.
230 |
231 | Parameters
232 | ----------
233 | filename : str
234 | Path to the mesh file.
235 |
236 | Returns
237 | -------
238 | meshes : list of :class:`~trimesh.base.Trimesh`
239 | The meshes loaded from the file.
240 | """
241 | meshes = trimesh.load(filename)
242 |
243 | # If we got a scene, dump the meshes
244 | if isinstance(meshes, trimesh.Scene):
245 | meshes = list(meshes.dump())
246 | meshes = [g for g in meshes if isinstance(g, trimesh.Trimesh)]
247 |
248 | if isinstance(meshes, (list, tuple, set)):
249 | meshes = list(meshes)
250 | if len(meshes) == 0:
251 | raise ValueError('At least one mesh must be pmeshesent in file')
252 | for r in meshes:
253 | if not isinstance(r, trimesh.Trimesh):
254 | raise TypeError('Could not load meshes from file')
255 | elif isinstance(meshes, trimesh.Trimesh):
256 | meshes = [meshes]
257 | else:
258 | raise ValueError('Unable to load mesh from file')
259 |
260 | return meshes
261 |
262 |
263 | def configure_origin(value):
264 | """Convert a value into a 4x4 transform matrix.
265 |
266 | Parameters
267 | ----------
268 | value : None, (6,) float, or (4,4) float
269 | The value to turn into the matrix.
270 | If (6,), interpreted as xyzrpy coordinates.
271 |
272 | Returns
273 | -------
274 | matrix : (4,4) float or None
275 | The created matrix.
276 | """
277 | if value is None:
278 | value = np.eye(4, dtype=np.float64)
279 | elif isinstance(value, (list, tuple, np.ndarray)):
280 | value = np.asanyarray(value, dtype=np.float64)
281 | if value.shape == (6,):
282 | value = xyz_rpy_to_matrix(value)
283 | elif value.shape != (4,4):
284 | raise ValueError('Origin must be specified as a 4x4 '
285 | 'homogenous transformation matrix')
286 | else:
287 | raise TypeError('Invalid type for origin, expect 4x4 matrix')
288 | return value
289 |
--------------------------------------------------------------------------------
/lib/dataset/const.py:
--------------------------------------------------------------------------------
1 | from .augmentations import (CropResizeToAspectAugmentation, PillowBlur,
2 | PillowBrightness, PillowColor, PillowContrast,
3 | PillowSharpness, occlusion_aug, to_torch_uint8)
4 |
5 | rgb_augmentations=[
6 | PillowSharpness(p=0.3, factor_interval=(0., 50.)),
7 | PillowContrast(p=0.3, factor_interval=(0.7, 1.8)),
8 | PillowBrightness(p=0.3, factor_interval=(0.7, 1.8)),
9 | PillowColor(p=0.3, factor_interval=(0., 4.))
10 | ]
11 |
12 | KEYPOINT_NAMES={
13 | 'panda' : [
14 | 'panda_link0', 'panda_link2', 'panda_link3',
15 | 'panda_link4', 'panda_link6', 'panda_link7',
16 | 'panda_hand'
17 | ],
18 | 'baxter': [
19 | 'torso_t0', 'right_s0','left_s0', 'right_s1', 'left_s1',
20 | 'right_e0','left_e0', 'right_e1','left_e1','right_w0', 'left_w0',
21 | 'right_w1','left_w1','right_w2', 'left_w2','right_hand','left_hand'
22 | ],
23 | 'kuka' : [
24 | 'iiwa7_link_0', 'iiwa7_link_1',
25 | 'iiwa7_link_2', 'iiwa7_link_3',
26 | 'iiwa7_link_4', 'iiwa7_link_5',
27 | 'iiwa7_link_6', 'iiwa7_link_7'
28 | ],
29 | 'owi535' :[
30 | 'Rotation', 'Base', 'Elbow', 'Wrist'
31 | ]
32 | }
33 |
34 | KEYPOINT_NAMES_TO_LINK_NAMES = {
35 | "panda" : dict(zip(KEYPOINT_NAMES['panda'],KEYPOINT_NAMES['panda'])),
36 | "kuka" : {
37 | 'iiwa7_link_0':'iiwa_link_0', 'iiwa7_link_1':'iiwa_link_1',
38 | 'iiwa7_link_2':'iiwa_link_2', 'iiwa7_link_3':'iiwa_link_3',
39 | 'iiwa7_link_4':'iiwa_link_4', 'iiwa7_link_5':'iiwa_link_5',
40 | 'iiwa7_link_6':'iiwa_link_6', 'iiwa7_link_7':'iiwa_link_7'
41 | },
42 | "baxter" : {
43 | 'torso_t0':'torso',
44 | 'right_s0':'right_upper_shoulder', 'left_s0':'left_upper_shoulder',
45 | 'right_s1':'right_lower_shoulder', 'left_s1':'left_lower_shoulder',
46 | 'right_e0':'right_upper_elbow','left_e0':'left_upper_elbow',
47 | 'right_e1':'right_lower_elbow','left_e1':'left_lower_elbow',
48 | 'right_w0':'right_upper_forearm', 'left_w0':'left_upper_forearm',
49 | 'right_w1':'right_lower_forearm', 'left_w1':'left_lower_forearm',
50 | 'right_w2':'right_wrist', 'left_w2':'left_wrist',
51 | 'right_hand':'right_hand','left_hand':'left_hand'
52 | },
53 | "owi535" : {
54 | 'Rotation':'Rotation', 'Base':'Base', 'Elbow':'Elbow', 'Wrist':'Wrist'
55 | }
56 | }
57 |
58 | LINK_NAMES = {
59 | 'panda': ['panda_link0', 'panda_link2', 'panda_link3', 'panda_link4',
60 | 'panda_link6', 'panda_link7', 'panda_hand'],
61 | 'kuka': ['iiwa_link_0', 'iiwa_link_1', 'iiwa_link_2', 'iiwa_link_3',
62 | 'iiwa_link_4', 'iiwa_link_5', 'iiwa_link_6', 'iiwa_link_7'],
63 | 'baxter': ['torso', 'right_upper_shoulder', 'left_upper_shoulder', 'right_lower_shoulder',
64 | 'left_lower_shoulder', 'right_upper_elbow', 'left_upper_elbow', 'right_lower_elbow',
65 | 'left_lower_elbow', 'right_upper_forearm', 'left_upper_forearm', 'right_lower_forearm',
66 | 'left_lower_forearm', 'right_wrist', 'left_wrist', 'right_hand', 'left_hand'],
67 | #'owi535': ["Base","Elbow","Wrist","Model","Model","Model","Model","Base","Base","Base","Base","Elbow","Elbow","Elbow","Elbow","Wrist","Wrist"],
68 | 'owi535' :[
69 | 'Rotation', 'Base', 'Elbow', 'Wrist'
70 | ]
71 | }
72 |
73 | JOINT_NAMES={
74 | 'panda': ['panda_joint1', 'panda_joint2', 'panda_joint3', 'panda_joint4',
75 | 'panda_joint5', 'panda_joint6', 'panda_joint7', 'panda_finger_joint1'],
76 | 'kuka': ['iiwa_joint_1', 'iiwa_joint_2', 'iiwa_joint_3', 'iiwa_joint_4',
77 | 'iiwa_joint_5', 'iiwa_joint_6', 'iiwa_joint_7'],
78 | 'baxter': ['head_pan', 'right_s0', 'left_s0', 'right_s1', 'left_s1',
79 | 'right_e0', 'left_e0', 'right_e1', 'left_e1', 'right_w0',
80 | 'left_w0', 'right_w1', 'left_w1', 'right_w2', 'left_w2'],
81 | 'owi535' :[
82 | 'Rotation', 'Base', 'Elbow', 'Wrist'
83 | ]
84 | }
85 |
86 | JOINT_TO_KP = {
87 | 'panda': [1, 1, 2, 3, 4, 4, 5, 6],
88 | 'kuka':[1,2,3,4,5,6,7],
89 | 'baxter':[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15],
90 | 'owi535':[0,1,2,3]
91 | }
92 |
93 | # flip_pairs=[
94 | # ["right_s0","left_s0"],["right_s1","left_s1"],["right_e0","left_e0"],
95 | # ["right_e1","left_e1"],["right_w0","left_w0"],["right_w1","left_w1"],
96 | # ["right_w2","left_w2"],["right_hand","left_hand"]
97 | # ]
98 | flip_pairs = [ [1,2],[3,4],[5,6],[7,8],[9,10],[11,12],[13,14],[15,16] ]
99 |
100 | PANDA_LIMB_LENGTH ={
101 | "link0-link2" : 0.3330,
102 | "link2-link3" : 0.3160,
103 | "link3-link4" : 0.0825,
104 | "link4-link6" : 0.39276,
105 | "link6-link7" : 0.0880,
106 | "link7-hand" : 0.1070
107 | }
108 | KUKA_LIMB_LENGTH ={
109 | "link0-link1" : 0.1500,
110 | "link1-link2" : 0.1900,
111 | "link2-link3" : 0.2100,
112 | "link3-link4" : 0.1900,
113 | "link4-link5" : 0.2100,
114 | "link5-link6" : 0.19946,
115 | "link6-link7" : 0.10122
116 | }
117 |
118 | LIMB_LENGTH = {
119 | "panda": list(PANDA_LIMB_LENGTH.values()),
120 | "kuka": list(KUKA_LIMB_LENGTH.values())
121 | }
122 |
123 | INITIAL_JOINT_ANGLE = {
124 | "zero": {
125 | "panda": {
126 | "panda_joint1": 0.0,
127 | "panda_joint2": 0.0,
128 | "panda_joint3": 0.0,
129 | "panda_joint4": 0.0,
130 | "panda_joint5": 0.0,
131 | "panda_joint6": 0.0,
132 | "panda_joint7": 0.0,
133 | "panda_finger_joint1": 0.0
134 | },
135 | "kuka": {
136 | "iiwa_joint_1": 0.0,
137 | "iiwa_joint_2": 0.0,
138 | "iiwa_joint_3": 0.0,
139 | "iiwa_joint_4": 0.0,
140 | "iiwa_joint_5": 0.0,
141 | "iiwa_joint_6": 0.0,
142 | "iiwa_joint_7": 0.0
143 | },
144 | "baxter": {
145 | "head_pan": 0.0,
146 | "right_s0": 0.0,
147 | "left_s0": 0.0,
148 | "right_s1": 0.0,
149 | "left_s1": 0.0,
150 | "right_e0": 0.0,
151 | "left_e0": 0.0,
152 | "right_e1": 0.0,
153 | "left_e1": 0.0,
154 | "right_w0": 0.0,
155 | "left_w0": 0.0,
156 | "right_w1": 0.0,
157 | "left_w1": 0.0,
158 | "right_w2": 0.0,
159 | "left_w2": 0.0
160 | },
161 | "owi535":{
162 | "Rotation":0.0,
163 | "Base":0.0,
164 | "Elbow":0.0,
165 | "Wrist":0.0
166 | }
167 | },
168 | "mean": {
169 | "panda": {
170 | "panda_joint1": 0.0,
171 | "panda_joint2": 0.0,
172 | "panda_joint3": 0.0,
173 | "panda_joint4": -1.52715,
174 | "panda_joint5": 0.0,
175 | "panda_joint6": 1.8675,
176 | "panda_joint7": 0.0,
177 | "panda_finger_joint1": 0.02
178 | },
179 | "kuka": {
180 | "iiwa_joint_1": 0.0,
181 | "iiwa_joint_2": 0.0,
182 | "iiwa_joint_3": 0.0,
183 | "iiwa_joint_4": 0.0,
184 | "iiwa_joint_5": 0.0,
185 | "iiwa_joint_6": 0.0,
186 | "iiwa_joint_7": 0.0
187 | },
188 | "baxter": {
189 | "head_pan": 0.0,
190 | "right_s0": 0.0,
191 | "left_s0": 0.0,
192 | "right_s1": -0.5499999999999999,
193 | "left_s1": -0.5499999999999999,
194 | "right_e0": 0.0,
195 | "left_e0": 0.0,
196 | "right_e1": 1.284,
197 | "left_e1": 1.284,
198 | "right_w0": 0.0,
199 | "left_w0": 0.0,
200 | "right_w1": 0.2616018366049999,
201 | "left_w1": 0.2616018366049999,
202 | "right_w2": 0.0,
203 | "left_w2": 0.0
204 | },
205 | "owi535":{
206 | "Rotation":0.0,
207 | "Base":-0.523598,
208 | "Elbow":0.523598,
209 | "Wrist":0.0
210 | }
211 | }
212 | }
213 |
214 | JOINT_BOUNDS = {
215 | "panda": [[-2.9671, 2.9671],
216 | [-1.8326, 1.8326],
217 | [-2.9671, 2.9671],
218 | [-3.1416, 0.0873],
219 | [-2.9671, 2.9671],
220 | [-0.0873, 3.8223],
221 | [-2.9671, 2.9671],
222 | [ 0.0000, 0.0400]],
223 |
224 | "kuka": [[-2.9671, 2.9671],
225 | [-2.0944, 2.0944],
226 | [-2.9671, 2.9671],
227 | [-2.0944, 2.0944],
228 | [-2.9671, 2.9671],
229 | [-2.0944, 2.0944],
230 | [-3.0543, 3.0543]],
231 |
232 | "baxter": [[-1.5708, 1.5708],
233 | [-1.7017, 1.7017],
234 | [-1.7017, 1.7017],
235 | [-2.1470, 1.0470],
236 | [-2.1470, 1.0470],
237 | [-3.0542, 3.0542],
238 | [-3.0542, 3.0542],
239 | [-0.0500, 2.6180],
240 | [-0.0500, 2.6180],
241 | [-3.0590, 3.0590],
242 | [-3.0590, 3.0590],
243 | [-1.5708, 2.0940],
244 | [-1.5708, 2.0940],
245 | [-3.0590, 3.0590],
246 | [-3.0590, 3.0590]],
247 | "owi535":[
248 | [-2.268928,2.268928],
249 | [-1.570796,1.047198],
250 | [-1.047198, 1.570796],
251 | [-0.785398,0.785398]
252 | ]
253 | }
254 |
255 |
256 | INTRINSICS_DICT = {
257 | "azure": (399.6578776041667, 399.4959309895833, 319.8955891927083, 244.0602823893229),
258 | "kinect": (525.0, 525.0, 319.5, 239.5),
259 | "realsense": (615.52392578125, 615.2191772460938, 328.2606506347656, 251.7917022705078),
260 | "orb": (615.52392578125, 615.2191772460938, 328.2606506347656, 251.7917022705078),
261 |
262 | }
263 |
--------------------------------------------------------------------------------
/lib/dataset/roboutils.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torchvision
3 | import numpy as np
4 | import torchvision.transforms as transforms
5 | from PIL import Image
6 | import sys
7 | sys.path.append("..")
8 | from utils.geometries import get_K_crop_resize
9 | import random
10 |
11 | def hnormalized(vector):
12 | hnormalized_vector = (vector / vector[-1])[:-1]
13 | return hnormalized_vector
14 |
15 | def crop_to_aspect_ratio(images, box, masks=None, K=None):
16 | assert images.dim() == 4
17 | bsz, _, h, w = images.shape
18 | assert box.dim() == 1
19 | assert box.shape[0] == 4
20 | w_output, h_output = box[[2, 3]] - box[[0, 1]]
21 | boxes = torch.cat(
22 | (torch.arange(bsz).unsqueeze(1).to(box.device).float(), box.unsqueeze(0).repeat(bsz, 1).float()),
23 | dim=1).to(images.device)
24 | images = torchvision.ops.roi_pool(images, boxes, output_size=(h_output, w_output))
25 | if masks is not None:
26 | assert masks.dim() == 4
27 | masks = torchvision.ops.roi_pool(masks, boxes, output_size=(h_output, w_output))
28 | if K is not None:
29 | assert K.dim() == 3
30 | assert K.shape[0] == bsz
31 | K = get_K_crop_resize(K, boxes[:, 1:], orig_size=(h, w), crop_resize=(h_output, w_output))
32 | return images, masks, K
33 |
34 |
35 | def make_detections_from_segmentation(masks):
36 | detections = []
37 | if masks.dim() == 4:
38 | assert masks.shape[0] == 1
39 | masks = masks.squeeze(0)
40 |
41 | for mask_n in masks:
42 | dets_n = dict()
43 | for uniq in torch.unique(mask_n, sorted=True):
44 | ids = np.where((mask_n == uniq).cpu().numpy())
45 | x1, y1, x2, y2 = np.min(ids[1]), np.min(ids[0]), np.max(ids[1]), np.max(ids[0])
46 | dets_n[int(uniq.item())] = torch.tensor([x1, y1, x2, y2]).to(mask_n.device)
47 | detections.append(dets_n)
48 | return detections
49 |
50 |
51 | def make_masks_from_det(detections, h, w):
52 | n_ids = len(detections)
53 | detections = torch.as_tensor(detections)
54 | masks = torch.zeros((n_ids, h, w)).byte()
55 | for mask_n, det_n in zip(masks, detections):
56 | x1, y1, x2, y2 = det_n.cpu().int().tolist()
57 | mask_n[y1:y2, x1:x2] = True
58 | return masks
59 |
60 | def get_bbox(bbox,w,h, strict=True):
61 | assert len(bbox)==4
62 | wmin, hmin, wmax, hmax = bbox
63 | if wmax<0 or hmax <0 or wmin > w or hmin > h:
64 | print("wmax",wmax,"hmax",hmax,"wmin",wmin,"hmin",hmin)
65 | wmin,hmin,wmax,hmax=max(0,wmin),max(0,hmin),min(w,wmax),min(h,hmax)
66 | wnew=wmax-wmin
67 | hnew=hmax-hmin
68 | wmin=int(max(0,wmin-0.3*wnew))
69 | wmax=int(min(w,wmax+0.3*wnew))
70 | hmin=int(max(0,hmin-0.3*hnew))
71 | hmax=int(min(h,hmax+0.3*hnew))
72 | wnew=wmax-wmin
73 | hnew=hmax-hmin
74 |
75 | if not strict:
76 | randomw = (random.random()-0.2)/2
77 | randomh = (random.random()-0.2)/2
78 |
79 | dwnew=randomw*wnew
80 | wmax+=dwnew/2
81 | wmin-=dwnew/2
82 |
83 | dhnew=randomh*hnew
84 | hmax+=dhnew/2
85 | hmin-=dhnew/2
86 |
87 | wmin=int(max(0,wmin))
88 | wmax=int(min(w,wmax))
89 | hmin=int(max(0,hmin))
90 | hmax=int(min(h,hmax))
91 | wnew=wmax-wmin
92 | hnew=hmax-hmin
93 |
94 | if wnew < 150:
95 | wmax+=75
96 | wmin-=75
97 | if hnew < 120:
98 | hmax+=60
99 | hmin-=60
100 |
101 | wmin,hmin,wmax,hmax=max(0,wmin),max(0,hmin),min(w,wmax),min(h,hmax)
102 | wmin,hmin,wmax,hmax=min(w,wmin),min(h,hmin),max(0,wmax),max(0,hmax)
103 | new_bbox = np.array([wmin,hmin,wmax,hmax])
104 | return new_bbox
105 |
106 | def get_bbox_raw(bbox):
107 | assert len(bbox)==4
108 | wmin, hmin, wmax, hmax = bbox
109 | wnew=wmax-wmin
110 | hnew=hmax-hmin
111 | wmin=int(wmin-0.3*wnew)
112 | wmax=int(wmax+0.3*wnew)
113 | hmin=int(hmin-0.3*hnew)
114 | hmax=int(hmax+0.3*hnew)
115 | wnew=wmax-wmin
116 | hnew=hmax-hmin
117 |
118 | if wnew < 150:
119 | wmax+=75
120 | wmin-=75
121 | if hnew < 120:
122 | hmax+=60
123 | hmin-=60
124 |
125 | new_bbox = np.array([wmin,hmin,wmax,hmax])
126 | return new_bbox
127 |
128 | def resize_image(image, bbox, mask, state, bbox_strict_bounded=None):
129 | #image as np.array
130 | wmin, hmin, wmax, hmax = bbox
131 | square_size =int(max(wmax - wmin, hmax - hmin))
132 | square_image = np.zeros((square_size, square_size, 3), dtype=np.uint8)
133 |
134 | x_offset = int((square_size - (wmax-wmin)) // 2)
135 | y_offset = int((square_size- (hmax-hmin)) // 2)
136 |
137 | square_image[y_offset:y_offset+(hmax-hmin), x_offset:x_offset+(wmax-wmin)] = image[hmin:hmax, wmin:wmax]
138 |
139 | keypoints=state['objects'][0]['keypoints_2d']
140 |
141 | for k in keypoints:
142 | k[1]-=hmin
143 | k[1]+=y_offset
144 | k[0]+=x_offset
145 | k[0]-=wmin
146 | if bbox_strict_bounded is not None:
147 | bbox_strict_bounded_new = bbox_strict_bounded[0]-wmin+x_offset, bbox_strict_bounded[1]-hmin+y_offset, \
148 | bbox_strict_bounded[2]-wmin+x_offset, bbox_strict_bounded[3]-hmin+y_offset
149 |
150 | K = state['camera']['K']
151 | K[0, 2] -= (wmin-x_offset)
152 | K[1, 2] -= (hmin-y_offset)
153 | if bbox_strict_bounded is None:
154 | return square_image, mask, state
155 | else:
156 | return square_image, mask, state, bbox_strict_bounded_new
157 |
158 | def tensor_to_image(tensor):
159 | image = tensor.cpu().clone().detach().numpy()
160 | image = Image.fromarray(image)
161 | return image
162 |
163 | def process_truncation(image, bbox, mask, state, max_pad=[120, 120, 120, 120]):
164 | #image as np.array
165 | wmin, hmin, wmax, hmax = bbox
166 | if wmin > 0 and hmin > 0 and hmax<480 and wmax <640:
167 | return image, bbox, mask, state
168 | d_wmin, d_hmin, d_wmax, d_hmax = int(-wmin), int(-hmin), int(wmax-640), int(hmax-480)
169 | d_wmin, d_hmin, d_wmax, d_hmax = int(max(0,d_wmin)), int(max(0,d_hmin)), int(max(0,d_wmax)), int(max(0,d_hmax))
170 | #print(d_wmin, d_hmin, d_wmax, d_hmax)
171 | d_wmin, d_hmin, d_wmax, d_hmax = min(max_pad[0],d_wmin), min(max_pad[1],d_hmin),min(max_pad[2],d_wmax),min(max_pad[3],d_hmax)
172 | wmax, hmax = 640 + d_wmax, 480+ d_hmax
173 | wnew, hnew = 640+d_wmax+d_wmin,480+d_hmax+d_hmin
174 |
175 | #print(wnew,hnew)
176 | new_image = np.zeros((hnew, wnew, 3), dtype=np.uint8)
177 |
178 | #print("d_hmin:",d_hmin,d_hmax, d_wmin, d_wmax,wnew, hnew,"hmax:",hmax)
179 | new_image[d_hmin:d_hmin+480, d_wmin:d_wmin+640] = image[0:480, 0:640]
180 |
181 |
182 | keypoints=state['objects'][0]['keypoints_2d']
183 |
184 | for k in keypoints:
185 | k[1]+=d_hmin
186 | k[0]+=d_wmin
187 |
188 | K = state['camera']['K']
189 | K[0, 2] += (d_wmin)
190 | K[1, 2] += (d_hmin)
191 |
192 | # new_bbox = np.array([max(0,int(wmin + d_wmin)),max(0,int(hmin + d_hmin)),int(wmax + d_wmin),int(hmax + d_hmin)])
193 | bbox_raw = np.concatenate([np.min(keypoints, axis=0)[0:2], np.max(keypoints, axis=0)[0:2]])
194 | new_bbox = get_bbox(bbox_raw,wnew,hnew)
195 | return new_image, new_bbox, mask, state
196 |
197 | def process_padding(image, bbox, mask, state, padding_pixel=25):
198 | #image as np.array
199 | keypoints=state['objects'][0]['keypoints_2d']
200 | # in_frame = 0
201 | # for k in keypoints:
202 | # if k[0]>0 and k[0]<256 and k[1]>0 and k[1]<256:
203 | # in_frame +=1
204 | # if in_frame ==7:
205 | # return image, bbox, mask, state
206 | # d_pad = 30 - 3*in_frame
207 | d_pad = padding_pixel
208 | d_wmin, d_hmin, d_wmax, d_hmax = d_pad,d_pad,d_pad,d_pad
209 |
210 | wnew, hnew = 320+d_wmax+d_wmin,320+d_hmax+d_hmin
211 |
212 | #print(wnew,hnew)
213 | new_image = np.zeros((hnew, wnew, 3), dtype=np.uint8)
214 |
215 | #print("d_hmin:",d_hmin,d_hmax, d_wmin, d_wmax,wnew, hnew,"hmax:",hmax)
216 | new_image[d_hmin:d_hmin+320, d_wmin:d_wmin+320] = image[0:320, 0:320]
217 |
218 | for k in keypoints:
219 | k[1]+=d_hmin
220 | k[0]+=d_wmin
221 |
222 | K = state['camera']['K']
223 | K[0, 2] += (d_wmin)
224 | K[1, 2] += (d_hmin)
225 |
226 | # new_bbox = np.array([max(0,int(wmin + d_wmin)),max(0,int(hmin + d_hmin)),int(wmax + d_wmin),int(hmax + d_hmin)])
227 | bbox_raw = np.concatenate([np.min(keypoints, axis=0)[0:2], np.max(keypoints, axis=0)[0:2]])
228 | new_bbox = get_bbox(bbox_raw,wnew,hnew)
229 | return new_image, new_bbox, mask, state
230 |
231 | def bbox_transform(bbox, K_original_inv, K, resize_hw):
232 | wmin, hmin, wmax, hmax = bbox
233 | corners = np.array([[wmin, hmin, 1.0],
234 | [wmax, hmin, 1.0],
235 | [wmax, hmax, 1.0],
236 | [wmin, hmax, 1.0]])
237 | corners3d_ill = np.matmul(K_original_inv, corners.T)
238 | new_corners = np.matmul(K, corners3d_ill).T
239 | assert all(new_corners[:,2] == 1.0), new_corners
240 | new_bbox = np.array([
241 | np.clip(new_corners[0,0], 0, resize_hw[0]),
242 | np.clip(new_corners[0,1], 0, resize_hw[1]),
243 | np.clip(new_corners[1,0], 0, resize_hw[0]),
244 | np.clip(new_corners[2,1], 0, resize_hw[1]),
245 | ])
246 | return new_bbox
247 |
248 | def get_extended_bbox(bbox, dwmin, dhmin, dwmax, dhmax, bounded=True, image_size=None):
249 | wmin, hmin, wmax, hmax = bbox
250 | extended_bbox = np.array([wmin-dwmin, hmin-dhmin, wmax+dwmax, hmax+dhmax])
251 | wmin, hmin, wmax, hmax = extended_bbox
252 | if bounded:
253 | assert image_size
254 | extended_bbox = np.array([max(0,wmin),max(0,hmin),min(image_size[0],wmax),min(image_size[1],hmax)])
255 | else:
256 | pass
257 | return extended_bbox
258 |
--------------------------------------------------------------------------------
/lib/utils/metrics.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import numpy as np
3 | from utils.transforms import point_projection_from_3d
4 | import matplotlib.pyplot as plt
5 | import seaborn as sns
6 | import os
7 |
8 | def compute_metrics_batch(robot,gt_keypoints3d,gt_keypoints2d,K_original,gt_joint,**pred_kwargs):
9 |
10 | # compute 3d keypoints locations
11 | # output shape: (batch_size, keypoints_num, 3)
12 | pred_joint = pred_kwargs["pred_joint"]
13 | pred_rot = pred_kwargs["pred_rot"]
14 | pred_trans = pred_kwargs["pred_trans"]
15 | if "pred_xy" in pred_kwargs and "pred_depth" in pred_kwargs and pred_kwargs["pred_xy"] is not None and pred_kwargs["pred_depth"] is not None:
16 | pred_xy = pred_kwargs["pred_xy"]
17 | pred_depth = pred_kwargs["pred_depth"]
18 | pred_trans = torch.cat((pred_xy,pred_depth),dim=-1)
19 | pred_xyz_integral = pred_kwargs["pred_xyz_integral"]
20 | reference_keypoint_id = pred_kwargs["reference_keypoint_id"]
21 |
22 | if pred_joint is None or pred_rot is None or pred_trans is None:
23 | assert pred_xyz_integral is not None
24 | pred_keypoints3d = pred_xyz_integral
25 | batch_size = pred_xyz_integral.shape[0]
26 | else:
27 | if reference_keypoint_id == 0:
28 | pred_keypoints3d = robot.get_keypoints(pred_joint,pred_rot,pred_trans)
29 | batch_size = pred_joint.shape[0]
30 | pred_joint = pred_joint.detach().cpu().numpy()
31 | else:
32 | pred_keypoints3d = robot.get_keypoints_root(pred_joint,pred_rot,pred_trans,root=reference_keypoint_id)
33 | batch_size = pred_joint.shape[0]
34 | pred_joint = pred_joint.detach().cpu().numpy()
35 |
36 | keypoints_num = len(robot.link_names)
37 | dof = robot.dof
38 | pred_keypoints3d = pred_keypoints3d.detach().cpu().numpy()
39 | gt_keypoints3d = gt_keypoints3d.detach().cpu().numpy()
40 | gt_keypoints2d = gt_keypoints2d.detach().cpu().numpy()
41 | K_original = K_original.detach().cpu().numpy()
42 | gt_joint = gt_joint.detach().cpu().numpy()
43 | pred_keypoints2d = point_projection_from_3d(K_original,pred_keypoints3d)
44 | assert(pred_keypoints3d.shape == (batch_size,keypoints_num,3)),f"{pred_keypoints3d.shape}"
45 | assert(gt_keypoints3d.shape == (batch_size,keypoints_num,3)),f"{gt_keypoints3d.shape}"
46 | assert(pred_keypoints2d.shape == (batch_size,keypoints_num,2)),f"{pred_keypoints2d.shape}"
47 | assert(gt_keypoints2d.shape == (batch_size,keypoints_num,2)),f"{gt_keypoints2d.shape}"
48 |
49 |
50 | # Thresholds (ADD:mm, PCK:pixel)
51 | add_thresholds = [1,5,10,20,40,60,80,100]
52 | pck_thresholds = [2.5,5.0,7.5,10.0,12.5,15.0,17.5,20.0]
53 |
54 | # ADD Average distance of detected keypoints within threshold
55 | error3d_batch = np.linalg.norm(pred_keypoints3d - gt_keypoints3d, ord = 2, axis = 2)
56 | assert(error3d_batch.shape == (batch_size,keypoints_num))
57 | error3d = np.mean(error3d_batch, axis = 1)
58 | # pcts3d = [len(np.where(error3d < th_mm/1000.0)[0])/float(error3d.shape[0]*error3d.shape[1]) for th_mm in add_thresholds]
59 |
60 | # PCK percentage of correct keypoints (only keypoints within the camera frame)
61 | error2d_batch = np.linalg.norm(pred_keypoints2d - gt_keypoints2d, ord = 2, axis = 2)
62 | assert(error2d_batch.shape == (batch_size,keypoints_num))
63 | valid = (gt_keypoints2d[:,:,0] <= 640.0) & (gt_keypoints2d[:,:,0] >= 0) & (gt_keypoints2d[:,:,1] <= 480.0) & (gt_keypoints2d[:,:,1] >= 0)
64 | error2d_all = error2d_batch * valid
65 | error2d_sum = np.sum(error2d_all, axis = 1)
66 | valid_sum = np.sum(valid, axis = 1)
67 | error2d = error2d_sum / valid_sum
68 | # pcts2d = [len(np.where(error2d < th_p)[0])/float(error2d.shape[0]*error2d.shape[1]) for th_p in pck_thresholds]
69 |
70 | # 3D/2D mean distance with gt of each keypoints
71 | dis3d = list(np.mean(error3d_batch, axis = 0))
72 | error2d_sum_batch = np.sum(error2d_all, axis = 0)
73 | valid_sum_batch = np.sum(valid, axis = 0)
74 | dis2d = error2d_sum_batch / valid_sum_batch
75 | # dis2d = list(np.mean(error2d_batch, axis = 0))
76 |
77 | # mean joint angle L1 error (per joint)
78 | # mean joint angle L1 error (per image)
79 | if pred_joint is not None:
80 | # pred_joint = pred_joint.detach().cpu().numpy()
81 | assert(gt_joint.shape == pred_joint.shape and gt_joint.shape == (batch_size, dof)), f"{pred_joint.shape},{gt_joint.shape}"
82 | error_joint = np.abs(gt_joint - pred_joint)
83 | l1_jointerror = list(np.mean(error_joint, axis = 0))
84 | if robot.robot_type == "panda":
85 | mean_jointerror = list(np.mean(error_joint[:,:-1], axis = 1))
86 | else:
87 | mean_jointerror = list(np.mean(error_joint, axis = 1))
88 | assert(len(mean_jointerror) == batch_size), len(mean_jointerror)
89 | else:
90 | l1_jointerror = [0] * dof
91 | mean_jointerror = [0] * batch_size
92 |
93 | # depth l1 error
94 | reference_keypoint_id = pred_kwargs["reference_keypoint_id"]
95 | error_depth = np.abs(pred_keypoints3d[:,reference_keypoint_id,2] - gt_keypoints3d[:,reference_keypoint_id,2])
96 |
97 | # root relative error
98 | pred_relatives = pred_keypoints3d[:,:,2] - pred_keypoints3d[:,reference_keypoint_id:reference_keypoint_id+1,2]
99 | gt_relatives = gt_keypoints3d[:,:,2] - gt_keypoints3d[:,reference_keypoint_id:reference_keypoint_id+1,2]
100 | error_relative = np.abs(pred_relatives - gt_relatives)
101 | batch_error_relative = np.mean(error_relative, axis=1)
102 |
103 | # root relative auc
104 | pred_keypoints3d_relative = pred_keypoints3d.copy()
105 | pred_keypoints3d_relative[:,:,2] = pred_relatives
106 | gt_keypoints3d_relative = gt_keypoints3d.copy()
107 | gt_keypoints3d_relative[:,:,2] = gt_relatives
108 | error3d_relative_batch = np.linalg.norm(pred_keypoints3d_relative - gt_keypoints3d_relative, ord = 2, axis = 2)
109 | assert(error3d_relative_batch.shape == (batch_size,keypoints_num))
110 | error3d_relative = np.mean(error3d_relative_batch, axis = 1)
111 |
112 |
113 |
114 | return error3d, error2d, dis3d, dis2d, l1_jointerror, mean_jointerror, error_depth, batch_error_relative, error3d_relative
115 |
116 |
117 | def summary_add_pck(alldis):
118 |
119 | dis3d = np.array(alldis['dis3d'])
120 | dis2d = np.array(alldis['dis2d'])
121 | assert(dis3d.shape[0] == dis2d.shape[0])
122 |
123 | add_threshold_ontb = [1,5,10,20,40,60,80,100]
124 | pck_threshold_ontb = [2.5,5.0,7.5,10.0,12.5,15.0,17.5,20.0]
125 |
126 | # for ADD
127 | auc_threshold = 0.1
128 | delta_threshold = 0.00001
129 | add_threshold_values = np.arange(0.0, auc_threshold, delta_threshold)
130 | counts_3d = []
131 | for value in add_threshold_values:
132 | under_threshold = (
133 | np.mean(dis3d <= value)
134 | )
135 | counts_3d.append(under_threshold)
136 | auc_add = np.trapz(counts_3d, dx=delta_threshold) / auc_threshold
137 |
138 | # for PCK
139 | auc_pixel_threshold = 20.0
140 | delta_pixel = 0.01
141 | pck_threshold_values = np.arange(0, auc_pixel_threshold, delta_pixel)
142 | counts_2d = []
143 | for value in pck_threshold_values:
144 | under_threshold = (
145 | np.mean(dis2d <= value)
146 | )
147 | counts_2d.append(under_threshold)
148 | auc_pck = np.trapz(counts_2d, dx=delta_pixel) / auc_pixel_threshold
149 |
150 | summary = {
151 | 'ADD/mean': np.mean(dis3d),
152 | 'ADD/median': np.median(dis3d),
153 | 'ADD/AUC': auc_add.item(),
154 | 'ADD_2D/mean': np.mean(dis2d),
155 | 'ADD_2D/median': np.median(dis2d),
156 | 'PCK/AUC': auc_pck.item()
157 | }
158 | for th_mm in add_threshold_ontb:
159 | summary[f'ADD_{th_mm}_mm'] = np.mean(dis3d <= th_mm * 1e-3)
160 | for th_p in pck_threshold_ontb:
161 | summary[f'PCK_{th_p}_pixel'] = np.mean(dis2d <= th_p)
162 | return summary
163 |
164 |
165 | def draw_add_curve(alldis, savename, testdsname, auc):
166 |
167 | dis3d = np.array(alldis['dis3d'])
168 | auc_threshold = 0.1
169 | delta_threshold = 0.00001
170 | add_threshold_values = np.arange(0.0, auc_threshold, delta_threshold)
171 | counts_3d = []
172 | for value in add_threshold_values:
173 | under_threshold = (
174 | np.mean(dis3d <= value)
175 | )
176 | counts_3d.append(under_threshold)
177 | plt.figure(figsize=(25,18))
178 | grid = plt.GridSpec(2,2, wspace=0.1, hspace=0.2)
179 | plt.subplot(grid[0,0])
180 | plt.grid()
181 | plt.plot(add_threshold_values, counts_3d)
182 | plt.xlim(0,auc_threshold)
183 | plt.ylim(0,1.0)
184 | plt.xlabel("add threshold values (unit: m)")
185 | plt.ylabel("percentages")
186 | plt.axvline(x=np.mean(dis3d), color='red', linestyle='--', label='mean distance')
187 | plt.axvline(x=np.median(dis3d), color='green', linestyle='--', label='median distance')
188 | plt.title("ADD curve")
189 | plt.text(x=0.001, y=0.9, s="auc="+str(round(auc*100, ndigits=2)))
190 | plt.legend()
191 |
192 | plt.subplot(grid[0,1])
193 | sns.histplot(dis3d, kde=True)
194 | plt.title("3d distance distribution, whole range")
195 |
196 | plt.subplot(grid[1,0])
197 | sns.histplot(dis3d, kde=True)
198 | plt.xlim(0, 0.5)
199 | plt.title("3d distance distribution, range: 0~0.5m")
200 |
201 | plt.subplot(grid[1,1])
202 | sns.histplot(dis3d, kde=True)
203 | plt.xlim(0, 0.1)
204 | plt.xticks(np.arange(0.0,0.101,0.01))
205 | plt.title("3d distance distribution, range: 0~0.1m")
206 | plt.axvline(x=np.mean(dis3d), color='red', linestyle='--', label='mean distance')
207 | plt.axvline(x=np.median(dis3d), color='green', linestyle='--', label='median distance')
208 |
209 | dataset_name = testdsname.split("/")[-1]
210 |
211 | plt.savefig(os.path.join(savename, f"add_distribution_curve_{dataset_name}.jpg"))
212 | print("drawn add curve in folder vis")
213 | plt.close()
214 |
215 |
216 | def draw_depth_figure(alldis, savename, testdsname):
217 | if "dr" in testdsname.split("/")[-1]:
218 | ds = "dr"
219 | elif "photo" in testdsname.split("/")[-1]:
220 | ds = "photo"
221 | else:
222 | ds = testdsname.split("/")[-1]
223 | assert len(alldis["deptherror"]) == len(alldis["gt_root_depth"]), (len(alldis["deptherror"]), len(alldis["gt_root_depth"]))
224 | deptherror = np.array(alldis["deptherror"])
225 | gtrootdepth = np.array(alldis["gt_root_depth"])
226 | plt.figure(figsize=(15,15))
227 | plt.scatter(gtrootdepth, deptherror)
228 | plt.xlim(0, 2.0)
229 | plt.ylim(0, 0.2)
230 | plt.title("root depth error -- gt root depth scatterplot")
231 | plt.savefig("unit_test/depth_curve/"+savename+"_"+ds+".jpg")
232 | plt.close()
233 |
234 | plt.close()
235 |
--------------------------------------------------------------------------------
/lib/utils/integral.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | from torch.nn import functional as F
4 | from utils.transforms import uvd_to_xyz
5 |
6 |
7 | def flip(x):
8 | assert (x.dim() == 3 or x.dim() == 4)
9 | dim = x.dim() - 1
10 |
11 | return x.flip(dims=(dim,))
12 |
13 | def norm_heatmap_hrnet(norm_type, heatmap, tau=5, sample_num=1):
14 | # Input tensor shape: [N,C,...]
15 | shape = heatmap.shape
16 | if norm_type == 'softmax':
17 | heatmap = heatmap.reshape(*shape[:2], -1)
18 | # global soft max
19 | heatmap = F.softmax(heatmap, 2)
20 | return heatmap.reshape(*shape)
21 | elif norm_type == 'sampling':
22 | heatmap = heatmap.reshape(*shape[:2], -1)
23 |
24 | eps = torch.rand_like(heatmap)
25 | log_eps = torch.log(-torch.log(eps))
26 | gumbel_heatmap = heatmap - log_eps / tau
27 |
28 | gumbel_heatmap = F.softmax(gumbel_heatmap, 2)
29 | return gumbel_heatmap.reshape(*shape)
30 | elif norm_type == 'multiple_sampling':
31 |
32 | heatmap = heatmap.reshape(*shape[:2], 1, -1)
33 |
34 | eps = torch.rand(*heatmap.shape[:2], sample_num, heatmap.shape[3], device=heatmap.device)
35 | log_eps = torch.log(-torch.log(eps))
36 | gumbel_heatmap = heatmap - log_eps / tau
37 | gumbel_heatmap = F.softmax(gumbel_heatmap, 3)
38 | gumbel_heatmap = gumbel_heatmap.reshape(shape[0], shape[1], sample_num, shape[2])
39 |
40 | # [B, S, K, -1]
41 | return gumbel_heatmap.transpose(1, 2)
42 | else:
43 | raise NotImplementedError
44 |
45 | def norm_heatmap_resnet(norm_type, heatmap):
46 | # Input tensor shape: [N,C,...]
47 | shape = heatmap.shape
48 | if norm_type == 'softmax':
49 | heatmap = heatmap.reshape(*shape[:2], -1)
50 | # global soft max
51 | heatmap = F.softmax(heatmap, 2)
52 | return heatmap.reshape(*shape)
53 | else:
54 | raise NotImplementedError
55 |
56 | def get_intrinsic_matrix_batch(f, c, bsz, inv=False):
57 |
58 | intrinsic_matrix = torch.zeros((bsz, 3, 3)).to(torch.float)
59 |
60 | if inv:
61 | intrinsic_matrix[:, 0, 0] = 1.0 / f[0].to(float)
62 | intrinsic_matrix[:, 0, 2] = - c[0].to(float) / f[0].to(float)
63 | intrinsic_matrix[:, 1, 1] = 1.0 / f[1].to(float)
64 | intrinsic_matrix[:, 1, 2] = - c[1].to(float) / f[1].to(float)
65 | intrinsic_matrix[:, 2, 2] = 1
66 | else:
67 | intrinsic_matrix[:, 0, 0] = f[0]
68 | intrinsic_matrix[:, 0, 2] = c[0]
69 | intrinsic_matrix[:, 1, 1] = f[1]
70 | intrinsic_matrix[:, 1, 2] = c[1]
71 | intrinsic_matrix[:, 2, 2] = 1
72 |
73 | return intrinsic_matrix.cuda(device=0)
74 |
75 | class HeatmapIntegralPose(nn.Module):
76 | """
77 | This module takes in heatmap output and performs soft-argmax(integral operation).
78 | """
79 | def __init__(self, backbone, **kwargs):
80 | super(HeatmapIntegralPose, self).__init__()
81 | self.backbone_name = backbone
82 | self.norm_type = kwargs["norm_type"]
83 | self.num_joints = kwargs["num_joints"]
84 | self.depth_dim = kwargs["depth_dim"]
85 | self.height_dim = kwargs["height_dim"]
86 | self.width_dim = kwargs["width_dim"]
87 | self.rootid = kwargs["rootid"] if "rootid" in kwargs else 0
88 | self.fixroot = kwargs["fixroot"] if "fixroot" in kwargs else False
89 |
90 | # self.focal_length = kwargs['FOCAL_LENGTH'] if 'FOCAL_LENGTH' in kwargs else 320
91 | bbox_3d_shape = kwargs['bbox_3d_shape'] if 'bbox_3d_shape' in kwargs else (2300, 2300, 2300)
92 | self.bbox_3d_shape = torch.tensor(bbox_3d_shape).float()
93 | self.depth_factor = self.bbox_3d_shape[2] * 1e-3
94 | self.image_size = kwargs["image_size"]
95 |
96 |
97 | def forward(self, out, flip_test=False, **kwargs):
98 | """
99 | Adapted from https://github.com/Jeff-sjtu/HybrIK/tree/main/hybrik/models
100 | """
101 |
102 | K = kwargs["K"]
103 | root_trans = kwargs["root_trans"]
104 | batch_size = out.shape[0]
105 | inv_k = get_intrinsic_matrix_batch((K[:,0,0],K[:,1,1]), (K[:,0,2],K[:,1,2]), bsz=batch_size, inv=True)
106 |
107 | if self.backbone_name in ["resnet", "resnet34", "resnet50"]:
108 | # out = out.reshape(batch_size, self.num_joints, self.depth_dim, self.height_dim, self.width_dim)
109 | out = out.reshape((out.shape[0], self.num_joints, -1))
110 | out = norm_heatmap_resnet(self.norm_type, out)
111 | assert out.dim() == 3, out.shape
112 | heatmaps = out / out.sum(dim=2, keepdim=True)
113 | heatmaps = heatmaps.reshape((heatmaps.shape[0], self.num_joints, self.depth_dim, self.height_dim, self.width_dim))
114 | hm_x0 = heatmaps.sum((2, 3)) # (B, K, W)
115 | hm_y0 = heatmaps.sum((2, 4)) # (B, K, H)
116 | hm_z0 = heatmaps.sum((3, 4)) # (B, K, D)
117 |
118 | range_tensor = torch.arange(hm_x0.shape[-1], dtype=torch.float32, device=hm_x0.device)
119 |
120 | hm_x = hm_x0 * range_tensor
121 | hm_y = hm_y0 * range_tensor
122 | hm_z = hm_z0 * range_tensor
123 |
124 | coord_x = hm_x.sum(dim=2, keepdim=True)
125 | coord_y = hm_y.sum(dim=2, keepdim=True)
126 | coord_z = hm_z.sum(dim=2, keepdim=True)
127 |
128 | coord_x = coord_x / float(self.width_dim) - 0.5
129 | coord_y = coord_y / float(self.height_dim) - 0.5
130 | coord_z = coord_z / float(self.depth_dim) - 0.5
131 |
132 | # -0.5 ~ 0.5
133 | pred_uvd_jts = torch.cat((coord_x, coord_y, coord_z), dim=2)
134 | if self.fixroot:
135 | pred_uvd_jts[:,self.rootid,2] = 0.0
136 | pred_uvd_jts_flat = pred_uvd_jts.reshape(batch_size, -1)
137 |
138 | pred_xyz_jts = uvd_to_xyz(uvd_jts=pred_uvd_jts, image_size=self.image_size, intrinsic_matrix_inverse=inv_k,
139 | root_trans=root_trans, depth_factor=self.depth_factor, return_relative=False)
140 |
141 | # pred_uvd_jts_back = xyz_to_uvd(xyz_jts=pred_xyz_jts, image_size=self.image_size, intrinsic_matrix=K,
142 | # root_trans=root_trans, depth_factor=self.depth_factor, return_relative=False)
143 | # print("(pred_uvd_jts-pred_uvd_jts_back).sum()",(pred_uvd_jts.cuda()-pred_uvd_jts_back.cuda()).sum())
144 |
145 | return pred_uvd_jts, pred_xyz_jts
146 |
147 | elif self.backbone_name == "hrnet" or self.backbone_name == "hrnet32" or self.backbone_name == "hrnet48":
148 | out = out.reshape((out.shape[0], self.num_joints, -1))
149 | heatmaps = norm_heatmap_hrnet(self.norm_type, out)
150 | assert heatmaps.dim() == 3, heatmaps.shape
151 | heatmaps = heatmaps.reshape((heatmaps.shape[0], self.num_joints, self.depth_dim, self.height_dim, self.width_dim))
152 |
153 | hm_x0 = heatmaps.sum((2, 3)) # (B, K, W)
154 | hm_y0 = heatmaps.sum((2, 4)) # (B, K, H)
155 | hm_z0 = heatmaps.sum((3, 4)) # (B, K, D)
156 |
157 | range_tensor = torch.arange(hm_x0.shape[-1], dtype=torch.float32, device=hm_x0.device).unsqueeze(-1)
158 | # hm_x = hm_x0 * range_tensor
159 | # hm_y = hm_y0 * range_tensor
160 | # hm_z = hm_z0 * range_tensor
161 |
162 | # coord_x = hm_x.sum(dim=2, keepdim=True)
163 | # coord_y = hm_y.sum(dim=2, keepdim=True)
164 | # coord_z = hm_z.sum(dim=2, keepdim=True)
165 | coord_x = hm_x0.matmul(range_tensor)
166 | coord_y = hm_y0.matmul(range_tensor)
167 | coord_z = hm_z0.matmul(range_tensor)
168 |
169 | coord_x = coord_x / float(self.width_dim) - 0.5
170 | coord_y = coord_y / float(self.height_dim) - 0.5
171 | coord_z = coord_z / float(self.depth_dim) - 0.5
172 |
173 | # -0.5 ~ 0.5
174 | pred_uvd_jts = torch.cat((coord_x, coord_y, coord_z), dim=2)
175 | if self.fixroot:
176 | pred_uvd_jts[:,self.rootid,2] = 0.0
177 | pred_uvd_jts_flat = pred_uvd_jts.reshape(batch_size, -1)
178 |
179 | pred_xyz_jts = uvd_to_xyz(uvd_jts=pred_uvd_jts, image_size=self.image_size, intrinsic_matrix_inverse=inv_k,
180 | root_trans=root_trans, depth_factor=self.depth_factor, return_relative=False)
181 |
182 | # pred_uvd_jts_back = xyz_to_uvd(xyz_jts=pred_xyz_jts, image_size=self.image_size, intrinsic_matrix=K,
183 | # root_trans=root_trans, depth_factor=self.depth_factor, return_relative=False)
184 | # print("(pred_uvd_jts-pred_uvd_jts_back).sum()",(pred_uvd_jts.cuda()-pred_uvd_jts_back.cuda()).sum())
185 |
186 | return pred_uvd_jts, pred_xyz_jts
187 |
188 | else:
189 | raise(NotImplementedError)
190 |
191 |
192 | class HeatmapIntegralJoint(nn.Module):
193 | """
194 | This module takes in heatmap output and performs soft-argmax(integral operation).
195 | """
196 | def __init__(self, backbone, **kwargs):
197 | super(HeatmapIntegralJoint, self).__init__()
198 | self.backbone_name = backbone
199 | self.norm_type = kwargs["norm_type"]
200 | self.dof = kwargs["dof"]
201 | self.joint_bounds = kwargs["joint_bounds"]
202 | assert self.joint_bounds.shape == (self.dof, 2), self.joint_bounds.shape
203 |
204 |
205 | def forward(self, out, **kwargs):
206 | """
207 | Adapted from https://github.com/Jeff-sjtu/HybrIK/tree/main/hybrik/models
208 | """
209 |
210 | batch_size = out.shape[0]
211 |
212 | if self.backbone_name in ["resnet34", "resnet50"]:
213 | out = out.reshape(batch_size, self.dof, -1)
214 | out = norm_heatmap_resnet(self.norm_type, out)
215 | assert out.dim() == 3, out.shape
216 | heatmaps = out / out.sum(dim=2, keepdim=True)
217 | heatmaps = heatmaps.reshape((heatmaps.shape[0], self.dof, -1)) # no depth dimension
218 |
219 | resolution = heatmaps.shape[-1]
220 | range_tensor = torch.arange(resolution, dtype=torch.float32, device=heatmaps.device).reshape(1,1,resolution)
221 | hm_int = heatmaps * range_tensor
222 | coord_joint_raw = hm_int.sum(dim=2, keepdim=True)
223 | coord_joint = coord_joint_raw / float(resolution) # 0~1
224 |
225 | bounds = self.joint_bounds.reshape(1,self.dof,2).cuda()
226 | jointrange = bounds[:,:,[1]] - bounds[:,:,[0]]
227 | joints = coord_joint * jointrange + bounds[:,:,[0]]
228 |
229 | return joints.squeeze(-1)
230 |
231 | else:
232 | raise(NotImplementedError)
233 |
--------------------------------------------------------------------------------
/lib/utils/utils.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
4 | import random
5 | import shutil
6 | from collections import OrderedDict, defaultdict
7 | from lib.dataset.multiepoch_dataloader import MultiEpochDataLoader
8 | from lib.dataset.samplers import PartialSampler
9 | from pathlib import Path
10 | import numpy as np
11 | import torch
12 | from lib.dataset.dream import DreamDataset
13 | from torch.utils.data import DataLoader
14 | from torch.utils.tensorboard import SummaryWriter
15 | from torchnet.meter import AverageValueMeter
16 | from tqdm import tqdm
17 |
18 | def cast(obj, device, dtype=None):
19 |
20 | if isinstance(obj, (dict, OrderedDict)):
21 | for k, v in obj.items():
22 | if v is None:
23 | continue
24 | obj[k] = cast(torch.as_tensor(v),device)
25 | if dtype is not None:
26 | obj[k] = obj[k].to(dtype)
27 | return obj
28 | else:
29 | return obj.to(device)
30 |
31 |
32 | def set_random_seed(seed):
33 |
34 | random.seed(seed)
35 | np.random.seed(seed)
36 | torch.manual_seed(seed)
37 | torch.cuda.manual_seed(seed)
38 |
39 |
40 | def copy_and_rename(src_path, dest_path, new_filename):
41 |
42 | src_path = Path(src_path)
43 | dest_path = Path(dest_path)
44 | shutil.copy(src_path, dest_path)
45 | src_filename = src_path.name
46 | dest_filepath = dest_path / new_filename
47 | (dest_path / src_filename).replace(dest_filepath)
48 |
49 |
50 | def create_logger(args):
51 |
52 | save_folder = os.path.join('experiments', args.exp_name)
53 | ckpt_folder = os.path.join(save_folder, 'ckpt')
54 | log_folder = os.path.join(save_folder, 'log')
55 | os.makedirs(ckpt_folder, exist_ok=True)
56 | os.makedirs(log_folder, exist_ok=True)
57 | writer = SummaryWriter(log_dir=log_folder)
58 | copy_and_rename(args.config_path, save_folder, "config.yaml")
59 |
60 | return save_folder, ckpt_folder, log_folder, writer
61 |
62 |
63 | def get_dataloaders(args):
64 |
65 | urdf_robot_name = args.urdf_robot_name
66 | train_ds_names = args.train_ds_names
67 | test_ds_name_dr = train_ds_names.replace("train_dr","test_dr")
68 | if urdf_robot_name != "baxter":
69 | test_ds_name_photo = train_ds_names.replace("train_dr","test_photo")
70 | if urdf_robot_name == "panda":
71 |
72 | test_ds_name_real = [train_ds_names.replace("synthetic/panda_synth_train_dr","real/panda-3cam_azure"),
73 | train_ds_names.replace("synthetic/panda_synth_train_dr","real/panda-3cam_kinect360"),
74 | train_ds_names.replace("synthetic/panda_synth_train_dr","real/panda-3cam_realsense"),
75 | train_ds_names.replace("synthetic/panda_synth_train_dr","real/panda-orb")]
76 |
77 | rootnet_hw = (int(args.rootnet_image_size),int(args.rootnet_image_size))
78 | other_hw = (int(args.other_image_size),int(args.other_image_size))
79 | ds_train = DreamDataset(train_ds_names,
80 | rootnet_resize_hw=rootnet_hw,
81 | other_resize_hw=other_hw,
82 | color_jitter=args.jitter, rgb_augmentation=args.other_aug,
83 | occlusion_augmentation=args.occlusion, occlu_p=args.occlu_p)
84 | ds_test_dr = DreamDataset(test_ds_name_dr,
85 | rootnet_resize_hw=rootnet_hw,
86 | other_resize_hw=other_hw,
87 | color_jitter=False, rgb_augmentation=False, occlusion_augmentation=False)
88 | if urdf_robot_name != "baxter":
89 | ds_test_photo = DreamDataset(test_ds_name_photo,
90 | rootnet_resize_hw=rootnet_hw,
91 | other_resize_hw=other_hw,
92 | color_jitter=False, rgb_augmentation=False, occlusion_augmentation=False)
93 |
94 | train_sampler = PartialSampler(ds_train, epoch_size=args.epoch_size)
95 | ds_iter_train = DataLoader(
96 | ds_train,
97 | sampler=train_sampler,
98 | batch_size=args.batch_size,
99 | num_workers=args.n_dataloader_workers,
100 | drop_last=False,
101 | pin_memory=True
102 | )
103 | ds_iter_train = MultiEpochDataLoader(ds_iter_train)
104 |
105 | test_loader_dict = {}
106 | ds_iter_test_dr = DataLoader(
107 | ds_test_dr,
108 | batch_size=args.batch_size,
109 | num_workers=args.n_dataloader_workers
110 | )
111 | test_loader_dict["dr"] = ds_iter_test_dr
112 |
113 | if urdf_robot_name != "baxter":
114 | ds_iter_test_photo = DataLoader(
115 | ds_test_photo,
116 | batch_size=args.batch_size,
117 | num_workers=args.n_dataloader_workers
118 | )
119 | test_loader_dict["photo"] = ds_iter_test_photo
120 |
121 | if urdf_robot_name == "panda":
122 | ds_shorts = ["azure", "kinect", "realsense", "orb"]
123 | for ds_name, ds_short in zip(test_ds_name_real, ds_shorts):
124 | ds_test_real = DreamDataset(ds_name,
125 | rootnet_resize_hw=rootnet_hw,
126 | other_resize_hw=other_hw,
127 | color_jitter=False, rgb_augmentation=False, occlusion_augmentation=False,
128 | process_truncation=args.fix_truncation)
129 | ds_iter_test_real = DataLoader(
130 | ds_test_real,
131 | batch_size=args.batch_size,
132 | num_workers=args.n_dataloader_workers
133 | )
134 | test_loader_dict[ds_short] = ds_iter_test_real
135 |
136 | print("len(ds_iter_train): ", len(ds_iter_train))
137 | print("len(ds_iter_test_dr): ", len(ds_iter_test_dr))
138 | if urdf_robot_name != "baxter":
139 | print("len(ds_iter_test_photo): ", len(ds_iter_test_photo))
140 | if urdf_robot_name == "panda":
141 | for ds_short in ds_shorts:
142 | print(f"len(ds_iter_test_{ds_short}): ", len(test_loader_dict[ds_short]))
143 |
144 | return ds_iter_train, test_loader_dict
145 |
146 |
147 | def get_scheduler(args, optimizer, last_epoch):
148 |
149 | def lr_lambda_linear(epoch):
150 | if epoch < args.n_epochs_warmup:
151 | ratio = float(epoch+1)/float(args.n_epochs_warmup)
152 | elif epoch <= args.start_decay:
153 | ratio = 1.0
154 | elif epoch <= args.end_decay:
155 | ratio = (float(args.end_decay - args.final_decay * args.start_decay) - (float(1-args.final_decay) * epoch)) / float(args.end_decay - args.start_decay)
156 | else:
157 | ratio = args.final_decay
158 | return ratio
159 |
160 | def lr_lambda_exponential(epoch):
161 | base_ratio = 1.0
162 | ratio = base_ratio
163 | if epoch < args.n_epochs_warmup:
164 | ratio = float(epoch+1)/float(args.n_epochs_warmup)
165 | elif epoch <= args.start_decay:
166 | ratio = base_ratio
167 | elif epoch <= args.end_decay:
168 | ratio = (args.exponent)**(epoch-args.start_decay)
169 | else:
170 | ratio = (args.exponent)**(args.end_decay-args.start_decay)
171 | return ratio
172 |
173 | def lr_lambda_everyXepoch(epoch):
174 | ratio = (args.step_decay)**(epoch // args.step)
175 | if epoch >= args.end_decay:
176 | ratio = (args.step_decay)**(args.end_decay // args.step)
177 | return ratio
178 |
179 | if args.use_schedule:
180 | if args.schedule_type == "linear":
181 | lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer=optimizer, lr_lambda=lr_lambda_linear, last_epoch=last_epoch)
182 | elif args.schedule_type == "exponential":
183 | lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer=optimizer, lr_lambda=lr_lambda_exponential, last_epoch=last_epoch)
184 | elif args.schedule_type == "everyXepoch":
185 | lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer=optimizer, lr_lambda=lr_lambda_everyXepoch, last_epoch=last_epoch)
186 | else:
187 | lr_scheduler = None
188 |
189 | return lr_scheduler
190 |
191 |
192 | def resume_run(args, model, optimizer, device):
193 |
194 | curr_max_auc_4real = { "azure": 0.0, "kinect": 0.0, "realsense": 0.0, "orb": 0.0 }
195 | template = 'ckpt/curr_best_auc(add)_DATASET_model.pk'
196 | ckpt_paths = [template.replace("DATASET", name) for name in curr_max_auc_4real.keys()]
197 |
198 | resume_dir = os.path.join("experiments" , args.resume_experiment_name)
199 | path = os.path.join(resume_dir, 'ckpt/curr_best_auc(add)_model.pk')
200 | checkpoint = torch.load(path)
201 | state_dict = checkpoint['model_state_dict']
202 | model.load_state_dict(state_dict)
203 | model.to(device)
204 | optimizer_dict = checkpoint['optimizer_state_dict']
205 | optimizer.load_state_dict(optimizer_dict)
206 | for state in optimizer.state.values():
207 | for k, v in state.items():
208 | if isinstance(v, torch.Tensor):
209 | state[k] = v.to(device)
210 |
211 | start_epoch = checkpoint['epoch']+1
212 | last_epoch = checkpoint['lr_scheduler_last_epoch']
213 | curr_max_auc = checkpoint["auc_add"]
214 |
215 | for postfix, dsname in zip(ckpt_paths, curr_max_auc_4real.keys()):
216 | model_path = os.path.join(resume_dir, postfix)
217 | ckpt = torch.load(model_path)
218 | curr_max_auc_onreal = ckpt["auc_add"]
219 | curr_max_auc_4real[dsname] = curr_max_auc_onreal
220 |
221 | return start_epoch, last_epoch, curr_max_auc, curr_max_auc_4real
222 |
223 |
224 | def save_checkpoint(args, auc_adds, model, optimizer, ckpt_folder, epoch, lr_scheduler, curr_max_auc, curr_max_auc_4real):
225 |
226 | save_path_dr = os.path.join(ckpt_folder, 'curr_best_auc(add)_model.pk')
227 | save_path_azure = os.path.join(ckpt_folder, 'curr_best_auc(add)_azure_model.pk')
228 | save_path_kinect = os.path.join(ckpt_folder, 'curr_best_auc(add)_kinect_model.pk')
229 | save_path_realsense = os.path.join(ckpt_folder, 'curr_best_auc(add)_realsense_model.pk')
230 | save_path_orb = os.path.join(ckpt_folder, 'curr_best_auc(add)_orb_model.pk')
231 | save_path = {"azure":save_path_azure, "kinect":save_path_kinect, "realsense":save_path_realsense, "orb":save_path_orb}
232 | saves = {"dr":True, "azure":True, "kinect":True, "realsense":True, "orb":True }
233 | if os.path.exists(save_path_dr):
234 | ckpt = torch.load(save_path_dr)
235 | if epoch <= ckpt["epoch"]: # prevent better model got covered during cluster rebooting
236 | saves["dr"] = False
237 | for real_name in ["azure", "kinect", "realsense", "orb"]:
238 | if os.path.exists(save_path[real_name]):
239 | ckpt_real = torch.load(save_path[real_name])
240 | if epoch <= ckpt_real["epoch"]: # prevent better model got covered during cluster rebooting
241 | saves[real_name] = False
242 |
243 | if saves["dr"]:
244 | if auc_adds["dr"] > curr_max_auc:
245 | curr_max_auc = auc_adds["dr"]
246 | last_epoch = lr_scheduler.last_epoch if args.use_schedule else -1
247 | torch.save({
248 | 'epoch': epoch,
249 | 'auc_add': curr_max_auc,
250 | 'model_state_dict': model.state_dict(),
251 | 'optimizer_state_dict': optimizer.state_dict(),
252 | 'lr_scheduler_last_epoch':last_epoch,
253 | }, save_path_dr)
254 |
255 | if args.urdf_robot_name == "panda":
256 | for real_name in ["azure", "kinect", "realsense", "orb"]:
257 | if saves[real_name]:
258 | if auc_adds[real_name] > curr_max_auc_4real[real_name]:
259 | curr_max_auc_4real[real_name] = auc_adds[real_name]
260 | last_epoch = lr_scheduler.last_epoch if args.use_schedule else -1
261 | torch.save({
262 | 'epoch': epoch,
263 | 'auc_add': curr_max_auc_4real[real_name],
264 | 'model_state_dict': model.state_dict(),
265 | 'optimizer_state_dict': optimizer.state_dict(),
266 | 'lr_scheduler_last_epoch':last_epoch,
267 | }, save_path[real_name])
--------------------------------------------------------------------------------
/lib/dataset/augmentations.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
4 | import math
5 | import random
6 | from copy import deepcopy
7 | import numpy as np
8 | import PIL
9 | import torch
10 | import torch.nn.functional as F
11 | from dataset.roboutils import hnormalized, make_detections_from_segmentation
12 | from PIL import ImageEnhance, ImageFilter
13 | from utils.geometries import get_K_crop_resize
14 |
15 |
16 | def to_pil(im):
17 | if isinstance(im, PIL.Image.Image):
18 | return im
19 | elif isinstance(im, torch.Tensor):
20 | return PIL.Image.fromarray(np.asarray(im))
21 | elif isinstance(im, np.ndarray):
22 | return PIL.Image.fromarray(im)
23 | else:
24 | raise ValueError('Type not supported', type(im))
25 |
26 |
27 | def to_torch_uint8(im):
28 | if isinstance(im, PIL.Image.Image):
29 | im = torch.as_tensor(np.asarray(im).astype(np.uint8))
30 | elif isinstance(im, torch.Tensor):
31 | assert im.dtype == torch.uint8
32 | elif isinstance(im, np.ndarray):
33 | assert im.dtype == np.uint8
34 | im = torch.as_tensor(im)
35 | else:
36 | raise ValueError('Type not supported', type(im))
37 | if im.dim() == 3:
38 | assert im.shape[-1] in {1, 3},f"{im.shape}"
39 | return im
40 |
41 | def occlusion_aug(bbox, img_shape, min_area=0.0, max_area=0.3, max_try_times=5):
42 | xmin, ymin, _, _ = bbox
43 | xmax = bbox[2]
44 | ymax = bbox[3]
45 | imght, imgwidth = img_shape
46 | counter = 0
47 | while True:
48 | # force to break if no suitable occlusion
49 | if counter > max_try_times: # 5
50 | return 0, 0, 0, 0
51 | counter += 1
52 |
53 | area_min = min_area # 0.0
54 | area_max = max_area # 0.3
55 | synth_area = (random.random() * (area_max - area_min) + area_min) * (xmax - xmin) * (ymax - ymin)
56 |
57 | ratio_min = 0.5
58 | ratio_max = 1 / 0.5
59 | synth_ratio = (random.random() * (ratio_max - ratio_min) + ratio_min)
60 |
61 | if(synth_ratio*synth_area<=0):
62 | print(synth_area,xmax,xmin,ymax,ymin)
63 | print(synth_ratio,ratio_max,ratio_min)
64 | synth_h = math.sqrt(synth_area * synth_ratio)
65 | synth_w = math.sqrt(synth_area / synth_ratio)
66 | synth_xmin = random.random() * ((xmax - xmin) - synth_w - 1) + xmin
67 | synth_ymin = random.random() * ((ymax - ymin) - synth_h - 1) + ymin
68 |
69 | if synth_xmin >= 0 and synth_ymin >= 0 and synth_xmin + synth_w < imgwidth and synth_ymin + synth_h < imght:
70 | synth_xmin = int(synth_xmin)
71 | synth_ymin = int(synth_ymin)
72 | synth_w = int(synth_w)
73 | synth_h = int(synth_h)
74 | break
75 | return synth_ymin, synth_h, synth_xmin, synth_w
76 |
77 | class PillowBlur:
78 | def __init__(self, p=0.4, factor_interval=(1, 3)):
79 | self.p = p
80 | self.factor_interval = factor_interval
81 |
82 | def __call__(self, im, mask, obs):
83 | im = to_pil(im)
84 | k = random.randint(*self.factor_interval)
85 | im = im.filter(ImageFilter.GaussianBlur(k))
86 | return im, mask, obs
87 |
88 |
89 | class PillowRGBAugmentation:
90 | def __init__(self, pillow_fn, p, factor_interval):
91 | self._pillow_fn = pillow_fn
92 | self.p = p
93 | self.factor_interval = factor_interval
94 |
95 | def __call__(self, im, mask, obs):
96 | im = to_pil(im)
97 | if random.random() <= self.p:
98 | im = self._pillow_fn(im).enhance(factor=random.uniform(*self.factor_interval))
99 | #im.save('./BRIGHT.png')
100 | return im, mask, obs
101 |
102 |
103 | class PillowSharpness(PillowRGBAugmentation):
104 | def __init__(self, p=0.3, factor_interval=(0., 50.)):
105 | super().__init__(pillow_fn=ImageEnhance.Sharpness,
106 | p=p,
107 | factor_interval=factor_interval)
108 |
109 |
110 | class PillowContrast(PillowRGBAugmentation):
111 | def __init__(self, p=0.3, factor_interval=(0.2, 50.)):
112 | super().__init__(pillow_fn=ImageEnhance.Contrast,
113 | p=p,
114 | factor_interval=factor_interval)
115 |
116 |
117 | class PillowBrightness(PillowRGBAugmentation):
118 | def __init__(self, p=0.5, factor_interval=(0.1, 6.0)):
119 | super().__init__(pillow_fn=ImageEnhance.Brightness,
120 | p=p,
121 | factor_interval=factor_interval)
122 |
123 |
124 | class PillowColor(PillowRGBAugmentation):
125 | def __init__(self, p=0.3, factor_interval=(0.0, 20.0)):
126 | super().__init__(pillow_fn=ImageEnhance.Color,
127 | p=p,
128 | factor_interval=factor_interval)
129 |
130 |
131 | class GrayScale(PillowRGBAugmentation):
132 | def __init__(self, p=0.3):
133 | self.p = p
134 |
135 | def __call__(self, im, mask, obs):
136 | im = to_pil(im)
137 | if random.random() <= self.p:
138 | im = to_torch_uint8(im).float()
139 | gray = 0.2989 * im[..., 0] + 0.5870 * im[..., 1] + 0.1140 * im[..., 2]
140 | gray = gray.to(torch.uint8)
141 | im = gray.unsqueeze(-1).repeat(1, 1, 3)
142 | return im, mask, obs
143 |
144 |
145 | class BackgroundAugmentation:
146 | def __init__(self, image_dataset, p):
147 | self.image_dataset = image_dataset
148 | self.p = p
149 |
150 | def get_bg_image(self, idx):
151 | return self.image_dataset[idx]
152 |
153 | def __call__(self, im, mask, obs):
154 | if random.random() <= self.p:
155 | im = to_torch_uint8(im)
156 | mask = to_torch_uint8(mask)
157 | h, w, c = im.shape
158 | im_bg = self.get_bg_image(random.randint(0, len(self.image_dataset) - 1))
159 | im_bg = to_pil(im_bg)
160 | im_bg = torch.as_tensor(np.asarray(im_bg.resize((w, h))))
161 | mask_bg = mask == 0
162 | im[mask_bg] = im_bg[mask_bg]
163 | return im, mask, obs
164 |
165 | class CropResizeToAspectAugmentation:
166 | def __init__(self, resize=(640, 480)):
167 | self.resize = (min(resize), max(resize))
168 | self.aspect = max(resize) / min(resize)
169 |
170 | def __call__(self, im, mask, obs, use_3d=True):
171 | im = to_torch_uint8(im)
172 | mask = to_torch_uint8(mask)
173 | obs['orig_camera'] = deepcopy(obs['camera'])
174 | assert im.shape[-1] == 3
175 | h, w = im.shape[:2]
176 | if (h, w) == self.resize:
177 | obs['orig_camera']['crop_resize_bbox'] = (0, 0, w-1, h-1)
178 | return im, mask, obs
179 |
180 | ratio = float(self.resize[0])/h
181 |
182 | images = (torch.as_tensor(im).float() / 255).unsqueeze(0).permute(0, 3, 1, 2)
183 | masks = torch.as_tensor(mask).unsqueeze(0).unsqueeze(0).float()
184 | K = torch.tensor(obs['camera']['K']).unsqueeze(0)
185 |
186 | # Match the width on input image with an image of target aspect ratio.
187 | # if not np.isclose(w/h, self.aspect):
188 | # x0, y0 = images.shape[-1] / 2, images.shape[-2] / 2
189 | # w = images.shape[-1]
190 | # r = self.aspect
191 | # h = w * 1/r
192 | # box_size = (h, w)
193 | # h, w = min(box_size), max(box_size)
194 | # x1, y1, x2, y2 = x0-w/2, y0-h/2, x0+w/2, y0+h/2
195 | # box = torch.tensor([x1, y1, x2, y2])
196 | # images, masks, K = crop_to_aspect_ratio(images, box, masks=masks, K=K)
197 |
198 | # Resize to target size
199 | x0, y0 = images.shape[-1] / 2, images.shape[-2] / 2
200 | h_input, w_input = images.shape[-2], images.shape[-1]
201 | h_output, w_output = min(self.resize), max(self.resize)
202 | box_size = (h_input, w_input)
203 | h, w = min(box_size), max(box_size)
204 | x1, y1, x2, y2 = x0-w/2, y0-h/2, x0+w/2, y0+h/2
205 | box = torch.tensor([x1, y1, x2, y2])
206 | images = F.interpolate(images, size=(h_output, w_output), mode='bilinear', align_corners=False)
207 | masks = F.interpolate(masks, size=(h_output, w_output), mode='nearest')
208 | obs['orig_camera']['crop_resize_bbox'] = tuple(box.tolist())
209 | K = get_K_crop_resize(K, box.unsqueeze(0), orig_size=(h_input, w_input), crop_resize=(h_output, w_output))
210 | # Update the bounding box annotations
211 | keypoints=[]
212 | keypoints_3d=obs['objects'][0]['TCO_keypoints_3d']
213 | K_tmp=K.cpu().clone().detach().numpy()[0]
214 |
215 | if use_3d:
216 | for location3d in keypoints_3d:
217 | location3d.reshape(3,1)
218 | p_unflattened = np.matmul(K_tmp, location3d)
219 | #print(p_unflattened)
220 | projection = hnormalized(p_unflattened)
221 | #print(p_unflattened)
222 | #print(projection)
223 | keypoints.append(list(projection))
224 | obs['objects'][0]['keypoints_2d']=keypoints
225 | else:
226 | obs['objects'][0]['keypoints_2d']=obs['objects'][0]['keypoints_2d']*ratio
227 |
228 | dets_gt = make_detections_from_segmentation(masks)[0]
229 | for n, obj in enumerate(obs['objects']):
230 | if 'bbox' in obj:
231 | #assert 'id_in_segm' in obj
232 | #print(dets_gt)
233 | try:
234 | obj['bbox'] = dets_gt[1]
235 | except:
236 | print("bbox",obj['bbox'],"dets_gt",dets_gt)
237 |
238 | im = (images[0].permute(1, 2, 0) * 255).to(torch.uint8)
239 | mask = masks[0, 0].to(torch.uint8)
240 | obs['camera']['K'] = K.squeeze(0).numpy()
241 | obs['camera']['resolution'] = (w_output, h_output)
242 | return im, mask, obs
243 |
244 | def flip_img(x):
245 | assert (x.ndim == 3)
246 | dim = 1
247 | return np.flip(x,axis=dim).copy()
248 |
249 | def flip_joints_2d(joints_2d, width, flip_pairs):
250 | joints = joints_2d.copy()
251 | joints[:, 0] = width - joints[:, 0] - 1 # flip horizontally
252 |
253 | if flip_pairs is not None: # change left-right parts
254 | for lr in flip_pairs:
255 | joints[lr[0]], joints[lr[1]] = joints[lr[1]].copy(), joints[lr[0]].copy()
256 | return joints
257 |
258 | def flip_xyz_joints_3d(joints_3d, flip_pairs):
259 | assert joints_3d.ndim in (2, 3)
260 | joints = joints_3d.copy()
261 | # flip horizontally
262 | joints[:, 0] = -1 * joints[:, 0]
263 | # change left-right parts
264 | if flip_pairs is not None:
265 | print(joints)
266 | for pair in flip_pairs:
267 | print(pair)
268 | print(joints[pair[0]], joints[pair[1]])
269 | joints[pair[0]], joints[pair[1]] = joints[pair[1]], joints[pair[0]].copy()
270 | return joints
271 |
272 | def flip_joints_3d(joints_3d, width, flip_pairs):
273 | joints = joints_3d.copy()
274 | # flip horizontally
275 | joints[:, 0, 0] = width - joints[:, 0, 0] - 1
276 | # change left-right parts
277 | if flip_pairs is not None:
278 | for pair in flip_pairs:
279 | joints[pair[0], :, 0], joints[pair[1], :, 0] = \
280 | joints[pair[1], :, 0], joints[pair[0], :, 0].copy()
281 | joints[pair[0], :, 1], joints[pair[1], :, 1] = \
282 | joints[pair[1], :, 1], joints[pair[0], :, 1].copy()
283 | joints[:, :, 0] *= joints[:, :, 1]
284 | return joints
285 |
286 | class FlipAugmentation:
287 | def __init__(self, p, flip_pairs=None):
288 | self.p = p
289 | self.flip_pairs = flip_pairs
290 |
291 | def __call__(self, im, mask, obs):
292 | if random.random() <= self.p:
293 | im = flip_img(im.numpy())
294 | # mask = flip_img(mask)
295 | obs['objects'][0]['keypoints_2d'] = flip_joints_2d(np.array(obs['objects'][0]['keypoints_2d']), im.shape[1], self.flip_pairs)
296 | obs['camera']['K'][0,0] = - obs['camera']['K'][0,0]
297 | obs['camera']['K'][0,2] = im.shape[1] - 1 - obs['camera']['K'][0,2]
298 | return im, mask, obs
299 |
300 | def rotate_joints_2d(joints_2d, width):
301 | joints = joints_2d.copy()
302 | joints[:, 1] = joints_2d[:, 0]
303 | joints[:, 0] = width - joints_2d[:, 1] + 1
304 | return joints
305 |
306 | class RotationAugmentation:
307 | def __init__(self, p):
308 | self.p = p
309 |
310 | def __call__(self, im, mask, obs):
311 | if random.random() <= self.p:
312 | h,w = im.shape[0],im.shape[1]
313 | im_copy = np.zeros((w,h,3), dtype=np.uint8)
314 | for i in range(h):
315 | for j in range(w):
316 | im_copy[j][h-i-1]=im[i][j].astype(np.uint8)
317 | rgb = PIL.fromarray(im_copy)
318 | obs['objects'][0]['keypoints_2d'] = rotate_joints_2d(obs['objects'][0]['keypoints_2d'], im_copy.shape[1])
319 | kp3d = obs['objects'][0]['TCO_keypoints_3d']
320 | K = obs['camera']['K']
321 | # original_fx,original_fy,original_cx,original_cy = K[0][0],K[1][1],K[0][2],K[1][2]
322 | # 角度（弧度）表示旋转方向，顺时针旋转90度
323 | angle = np.pi / 2 # 90 degree
324 | K[0][2],K[1][2] = K[1][2],K[0][2]
325 | # set up rotation matrix
326 | rotation_matrix = np.array([[np.cos(angle), -np.sin(angle), 0],
327 | [np.sin(angle), np.cos(angle), 0],
328 | [0, 0, 1]])
329 | # new camera intrinsic matrix
330 | # obs['camera']['K'] = new_intrinsic_matrix
331 | for i in range(kp3d.shape[0]):
332 | kp3d[i] = np.dot(rotation_matrix, kp3d[i])
333 |
334 | return rgb, mask, obs
335 | else:
336 | pass
--------------------------------------------------------------------------------
/lib/utils/geometries.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import torch
3 | from torch.nn import functional as F
4 |
5 | def batch_rodrigues(theta):
6 | """Convert axis-angle representation to rotation matrix.
7 | Args:
8 | theta: size = [B, 3]
9 | Returns:
10 | Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
11 | """
12 | l1norm = torch.norm(theta + 1e-8, p = 2, dim = 1)
13 | angle = torch.unsqueeze(l1norm, -1)
14 | normalized = torch.div(theta, angle)
15 | angle = angle * 0.5
16 | v_cos = torch.cos(angle)
17 | v_sin = torch.sin(angle)
18 | quat = torch.cat([v_cos, v_sin * normalized], dim = 1)
19 | return quat_to_rotmat(quat)
20 |
21 | def quat_to_rotmat(quat):
22 | """Convert quaternion coefficients to rotation matrix.
23 | Args:
24 | quat: size = [B, 4] 4 <===>(w, x, y, z)
25 | Returns:
26 | Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
27 | """
28 | norm_quat = quat
29 | norm_quat = norm_quat/(norm_quat.norm(p=2, dim=1, keepdim=True)+1e-9)
30 | w, x, y, z = norm_quat[:,0], norm_quat[:,1], norm_quat[:,2], norm_quat[:,3]
31 |
32 | B = quat.size(0)
33 |
34 | w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)
35 | wx, wy, wz = w*x, w*y, w*z
36 | xy, xz, yz = x*y, x*z, y*z
37 |
38 | rotMat = torch.stack([w2 + x2 - y2 - z2, 2*xy - 2*wz, 2*wy + 2*xz,
39 | 2*wz + 2*xy, w2 - x2 + y2 - z2, 2*yz - 2*wx,
40 | 2*xz - 2*wy, 2*wx + 2*yz, w2 - x2 - y2 + z2], dim=1).view(B, 3, 3)
41 | return rotMat
42 |
43 | def quat_to_rotmat_np(quat):
44 | """Convert quaternion coefficients to rotation matrix.
45 | Without batch, in numpy.
46 | Args:
47 | quat: size = [4] <===>(w, x, y, z)
48 | Returns:
49 | Rotation matrix corresponding to the quaternion -- size = [3, 3]
50 | """
51 | norm_quat = quat
52 | norm_quat = norm_quat / np.linalg.norm(norm_quat, ord=2, axis=0, keepdims=True)
53 | w, x, y, z = norm_quat[0], norm_quat[1], norm_quat[2], norm_quat[3]
54 | w2, x2, y2, z2 = w*w, x*x, y*y, z*z
55 | wx, wy, wz = w*x, w*y, w*z
56 | xy, xz, yz = x*y, x*z, y*z
57 |
58 | rotMat = np.array([[w2 - x2 - y2 + z2, -2*yz + 2*wx, 2*wy + 2*xz],
59 | [2*wx + 2*yz, -(w2 - x2 + y2 - z2), 2*xy - 2*wz],
60 | [-2*xz + 2*wy, 2*wz + 2*xy, -(w2 + x2 - y2 - z2)]])
61 | return rotMat
62 |
63 | def rotmat_to_quat(matrices):
64 | batch = matrices.shape[0]
65 | this_device = matrices.device
66 | w = torch.sqrt(torch.max(1.0 + matrices[:,0,0] + matrices[:,1,1] + matrices[:,2,2], torch.zeros(1).to(this_device))) / 2.0
67 | w = torch.max (w , torch.autograd.Variable(torch.zeros(batch)).to(this_device) + 1e-8) #batch
68 | w4 = 4.0 * w
69 | x = (matrices[:,2,1] - matrices[:,1,2]) / w4
70 | y = (matrices[:,0,2] - matrices[:,2,0]) / w4
71 | z = (matrices[:,1,0] - matrices[:,0,1]) / w4
72 | quats = torch.cat( (w.view(batch,1), x.view(batch, 1),y.view(batch, 1), z.view(batch, 1) ), 1 )
73 | quats = normalize_vector(quats)
74 | return quats
75 |
76 | def normalize_vector(v):
77 | batch = v.shape[0]
78 | v_mag = torch.sqrt(v.pow(2).sum(1))# batch
79 | v_mag = torch.max(v_mag, torch.autograd.Variable(torch.FloatTensor([1e-8]).to(v.device)))
80 | v_mag = v_mag.view(batch,1).expand(batch,v.shape[1])
81 | v = v/v_mag
82 | return v
83 |
84 | # def rot6d_to_rotmat(x):
85 | # """Convert 6D rotation representation to 3x3 rotation matrix.
86 | # Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019
87 | # Input:
88 | # (B,6) Batch of 6-D rotation representations
89 | # Output:
90 | # (B,3,3) Batch of corresponding rotation matrices
91 | # """
92 | # x = x.view(-1,3,2)
93 | # a1 = x[:, :, 0]
94 | # a2 = x[:, :, 1]
95 | # b1 = F.normalize(a1)
96 | # b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1)
97 | # b3 = torch.cross(b1, b2)
98 | # return torch.stack((b1, b2, b3), dim=-1)
99 |
100 | def rot6d_to_rotmat(poses):
101 | """
102 | Code from https://github.com/papagina/RotationContinuity
103 | On the Continuity of Rotation Representations in Neural Networks
104 | Zhou et al. CVPR19
105 | https://zhouyisjtu.github.io/project_rotation/rotation.html
106 | """
107 | assert poses.shape[-1] == 6
108 | x_raw = poses[..., 0:3]
109 | y_raw = poses[..., 3:6]
110 | x = x_raw / torch.norm(x_raw, p=2, dim=-1, keepdim=True)
111 | z = torch.cross(x, y_raw, dim=-1)
112 | z = z / torch.norm(z, p=2, dim=-1, keepdim=True)
113 | y = torch.cross(z, x, dim=-1)
114 | matrix = torch.stack((x, y, z), -1)
115 | return torch.transpose(matrix,dim0=-2,dim1=-1)
116 |
117 | def rotmat_to_rot6d(matrix):
118 | """
119 | Converts rotation matrices to 6D rotation representation by Zhou et al. [1]
120 | by dropping the last row. Note that 6D representation is not unique.
121 | Args:
122 | matrix: batch of rotation matrices of size (*, 3, 3)
123 |
124 | Returns:
125 | 6D rotation representation, of size (*, 6)
126 |
127 | [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
128 | On the Continuity of Rotation Representations in Neural Networks.
129 | IEEE Conference on Computer Vision and Pattern Recognition, 2019.
130 | Retrieved from http://arxiv.org/abs/1812.07035
131 | """
132 | return matrix[..., :2, :].clone().reshape(*matrix.size()[:-2], 6)
133 |
134 | def rot9d_to_rotmat(x):
135 | """
136 | Maps 9D input vectors onto SO(3) via symmetric orthogonalization.
137 | x: should have size [batch_size, 9]
138 | Output has size [batch_size, 3, 3], where each inner 3x3 matrix is in SO(3).
139 | """
140 | m = x.view(-1, 3, 3)
141 | d = m.device
142 | u, s, v = torch.svd(m.cpu())
143 | u, v = u.to(d), v.to(d)
144 | vt = torch.transpose(v, 1, 2)
145 | det = torch.det(torch.bmm(u, vt))
146 | det = det.view(-1, 1, 1)
147 | vt = torch.cat((vt[:, :2, :], vt[:, -1:, :] * det), 1)
148 | r = torch.bmm(u, vt)
149 | return r.cuda()
150 |
151 | #matrices batch*3*3
152 | #both matrix are orthogonal rotation matrices
153 | #out theta between 0 to 180 degree batch
154 | def compute_geodesic_distance_from_two_matrices(m1, m2):
155 | batch=m1.shape[0]
156 | m = torch.bmm(m1, m2.transpose(1,2)) #batch*3*3
157 | cos = ( m[:,0,0] + m[:,1,1] + m[:,2,2] - 1 )/2
158 | cos = torch.min(cos, torch.autograd.Variable(torch.ones(batch).cuda()) )
159 | cos = torch.max(cos, torch.autograd.Variable(torch.ones(batch).cuda())*-1 )
160 | theta = torch.acos(cos)
161 | # theta = torch.min(theta, 2*np.pi - theta)
162 | return theta
163 |
164 | def angle_axis_to_rotation_matrix(angle_axis):
165 | """Convert 3d vector of axis-angle rotation to 4x4 rotation matrix
166 |
167 | Args:
168 | angle_axis (Tensor): tensor of 3d vector of axis-angle rotations.
169 |
170 | Returns:
171 | Tensor: tensor of 4x4 rotation matrices.
172 |
173 | Shape:
174 | - Input: :math:`(N, 3)`
175 | - Output: :math:`(N, 4, 4)`
176 |
177 | Example:
178 | >>> input = torch.rand(1, 3) # Nx3
179 | >>> output = tgm.angle_axis_to_rotation_matrix(input) # Nx4x4
180 | """
181 | def _compute_rotation_matrix(angle_axis, theta2, eps=1e-6):
182 | # We want to be careful to only evaluate the square root if the
183 | # norm of the angle_axis vector is greater than zero. Otherwise
184 | # we get a division by zero.
185 | k_one = 1.0
186 | theta = torch.sqrt(theta2)
187 | wxyz = angle_axis / (theta + eps)
188 | wx, wy, wz = torch.chunk(wxyz, 3, dim=1)
189 | cos_theta = torch.cos(theta)
190 | sin_theta = torch.sin(theta)
191 |
192 | r00 = cos_theta + wx * wx * (k_one - cos_theta)
193 | r10 = wz * sin_theta + wx * wy * (k_one - cos_theta)
194 | r20 = -wy * sin_theta + wx * wz * (k_one - cos_theta)
195 | r01 = wx * wy * (k_one - cos_theta) - wz * sin_theta
196 | r11 = cos_theta + wy * wy * (k_one - cos_theta)
197 | r21 = wx * sin_theta + wy * wz * (k_one - cos_theta)
198 | r02 = wy * sin_theta + wx * wz * (k_one - cos_theta)
199 | r12 = -wx * sin_theta + wy * wz * (k_one - cos_theta)
200 | r22 = cos_theta + wz * wz * (k_one - cos_theta)
201 | rotation_matrix = torch.cat(
202 | [r00, r01, r02, r10, r11, r12, r20, r21, r22], dim=1)
203 | return rotation_matrix.view(-1, 3, 3)
204 |
205 | def _compute_rotation_matrix_taylor(angle_axis):
206 | rx, ry, rz = torch.chunk(angle_axis, 3, dim=1)
207 | k_one = torch.ones_like(rx)
208 | rotation_matrix = torch.cat(
209 | [k_one, -rz, ry, rz, k_one, -rx, -ry, rx, k_one], dim=1)
210 | return rotation_matrix.view(-1, 3, 3)
211 |
212 | # stolen from ceres/rotation.h
213 |
214 | _angle_axis = torch.unsqueeze(angle_axis, dim=1)
215 | theta2 = torch.matmul(_angle_axis, _angle_axis.transpose(1, 2))
216 | theta2 = torch.squeeze(theta2, dim=1)
217 |
218 | # compute rotation matrices
219 | rotation_matrix_normal = _compute_rotation_matrix(angle_axis, theta2)
220 | rotation_matrix_taylor = _compute_rotation_matrix_taylor(angle_axis)
221 |
222 | # create mask to handle both cases
223 | eps = 1e-6
224 | mask = (theta2 > eps).view(-1, 1, 1).to(theta2.device)
225 | mask_pos = (mask).type_as(theta2)
226 | mask_neg = (mask == False).type_as(theta2) # noqa
227 |
228 | # create output pose matrix
229 | batch_size = angle_axis.shape[0]
230 | rotation_matrix = torch.eye(4).to(angle_axis.device).type_as(angle_axis)
231 | rotation_matrix = rotation_matrix.view(1, 4, 4).repeat(batch_size, 1, 1)
232 | # fill output matrix with masked values
233 | rotation_matrix[..., :3, :3] = \
234 | mask_pos * rotation_matrix_normal + mask_neg * rotation_matrix_taylor
235 | return rotation_matrix # Nx4x4
236 |
237 |
238 | def perspective_projection(points, rotation, translation,
239 | focal_length, camera_center):
240 | """
241 | This function computes the perspective projection of a set of points.
242 | Input:
243 | points (bs, N, 3): 3D points
244 | rotation (bs, 3, 3): Camera rotation
245 | translation (bs, 3): Camera translation
246 | focal_length (bs,) or scalar: Focal length
247 | camera_center (bs, 2): Camera center
248 | """
249 | batch_size = points.shape[0]
250 | K = torch.zeros([batch_size, 3, 3], device=points.device)
251 | K[:,0,0] = focal_length
252 | K[:,1,1] = focal_length
253 | K[:,2,2] = 1.
254 | K[:,:-1, -1] = camera_center
255 |
256 | # Transform points
257 | points = torch.einsum('bij,bkj->bki', rotation, points)
258 | points = points + translation.unsqueeze(1)
259 |
260 | # Apply perspective distortion
261 | projected_points = points / points[:,:,-1].unsqueeze(-1)
262 |
263 | # Apply camera intrinsics
264 | projected_points = torch.einsum('bij,bkj->bki', K, projected_points)
265 |
266 | return projected_points[:, :, :-1]
267 |
268 |
269 | def estimate_translation_np(S, joints_2d, joints_conf, focal_length=5000, img_size=224):
270 | """Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d.
271 | Input:
272 | S: (25, 3) 3D joint locations
273 | joints: (25, 3) 2D joint locations and confidence
274 | Returns:
275 | (3,) camera translation vector
276 | """
277 |
278 | num_joints = S.shape[0]
279 | # focal length
280 | f = np.array([focal_length,focal_length])
281 | # optical center
282 | center = np.array([img_size/2., img_size/2.])
283 |
284 | # transformations
285 | Z = np.reshape(np.tile(S[:,2],(2,1)).T,-1)
286 | XY = np.reshape(S[:,0:2],-1)
287 | O = np.tile(center,num_joints)
288 | F = np.tile(f,num_joints)
289 | weight2 = np.reshape(np.tile(np.sqrt(joints_conf),(2,1)).T,-1)
290 |
291 | # least squares
292 | Q = np.array([F*np.tile(np.array([1,0]),num_joints), F*np.tile(np.array([0,1]),num_joints), O-np.reshape(joints_2d,-1)]).T
293 | c = (np.reshape(joints_2d,-1)-O)*Z - F*XY
294 |
295 | # weighted least squares
296 | W = np.diagflat(weight2)
297 | Q = np.dot(W,Q)
298 | c = np.dot(W,c)
299 |
300 | # square matrix
301 | A = np.dot(Q.T,Q)
302 | b = np.dot(Q.T,c)
303 |
304 | # solution
305 | trans = np.linalg.solve(A, b)
306 |
307 | return trans
308 |
309 |
310 | def estimate_translation(S, joints_2d, focal_length=5000., img_size=224.):
311 | """Find camera translation that brings 3D joints S closest to 2D the corresponding joints_2d.
312 | Input:
313 | S: (B, 49, 3) 3D joint locations
314 | joints: (B, 49, 3) 2D joint locations and confidence
315 | Returns:
316 | (B, 3) camera translation vectors
317 | """
318 |
319 | device = S.device
320 | # Use only joints 25:49 (GT joints)
321 | S = S[:, 25:, :].cpu().numpy()
322 | joints_2d = joints_2d[:, 25:, :].cpu().numpy()
323 | joints_conf = joints_2d[:, :, -1]
324 | joints_2d = joints_2d[:, :, :-1]
325 | trans = np.zeros((S.shape[0], 3), dtype=np.float32)
326 | # Find the translation for each example in the batch
327 | for i in range(S.shape[0]):
328 | S_i = S[i]
329 | joints_i = joints_2d[i]
330 | conf_i = joints_conf[i]
331 | trans[i] = estimate_translation_np(S_i, joints_i, conf_i, focal_length=focal_length, img_size=img_size)
332 | return torch.from_numpy(trans).to(device)
333 |
334 | #input batch*4*4 or batch*3*3
335 | #output torch batch*3 x, y, z in radiant
336 | #the rotation is in the sequence of x,y,z
337 | def compute_euler_angles_from_rotation_matrices(rotation_matrices):
338 | batch=rotation_matrices.shape[0]
339 | R=rotation_matrices
340 | sy = torch.sqrt(R[:,0,0]*R[:,0,0]+R[:,1,0]*R[:,1,0])
341 | singular= sy<1e-6
342 | singular=singular.float()
343 |
344 | x=torch.atan2(R[:,2,1], R[:,2,2])
345 | y=torch.atan2(-R[:,2,0], sy)
346 | z=torch.atan2(R[:,1,0],R[:,0,0])
347 |
348 | xs=torch.atan2(-R[:,1,2], R[:,1,1])
349 | ys=torch.atan2(-R[:,2,0], sy)
350 | zs=R[:,1,0]*0
351 |
352 | out_euler=torch.autograd.Variable(torch.zeros(batch,3).cuda())
353 | out_euler[:,0]=x*(1-singular)+xs*singular
354 | out_euler[:,1]=y*(1-singular)+ys*singular
355 | out_euler[:,2]=z*(1-singular)+zs*singular
356 |
357 | return out_euler
358 |
359 |
360 | def get_K_crop_resize(K, boxes, orig_size, crop_resize):
361 | """
362 | Adapted from https://github.com/BerkeleyAutomation/perception/blob/master/perception/camera_intrinsics.py
363 | Skew is not handled !
364 | """
365 | assert K.shape[1:] == (3, 3)
366 | assert boxes.shape[1:] == (4, )
367 | K = K.float()
368 | boxes = boxes.float()
369 | new_K = K.clone()
370 |
371 | orig_size = torch.tensor(orig_size, dtype=torch.float)
372 | crop_resize = torch.tensor(crop_resize, dtype=torch.float)
373 |
374 | final_width, final_height = max(crop_resize), min(crop_resize)
375 | crop_width = boxes[:, 2] - boxes[:, 0]
376 | crop_height = boxes[:, 3] - boxes[:, 1]
377 | crop_cj = (boxes[:, 0] + boxes[:, 2]) / 2
378 | crop_ci = (boxes[:, 1] + boxes[:, 3]) / 2
379 |
380 | # Crop
381 | cx = K[:, 0, 2] + (crop_width - 1) / 2 - crop_cj
382 | cy = K[:, 1, 2] + (crop_height - 1) / 2 - crop_ci
383 |
384 | # # Resize (upsample)
385 | center_x = (crop_width - 1) / 2
386 | center_y = (crop_height - 1) / 2
387 | orig_cx_diff = cx - center_x
388 | orig_cy_diff = cy - center_y
389 | scale_x = final_width / crop_width
390 | scale_y = final_height / crop_height
391 | scaled_center_x = (final_width - 1) / 2
392 | scaled_center_y = (final_height - 1) / 2
393 | fx = scale_x * K[:, 0, 0]
394 | fy = scale_y * K[:, 1, 1]
395 | cx = scaled_center_x + scale_x * orig_cx_diff
396 | cy = scaled_center_y + scale_y * orig_cy_diff
397 |
398 | new_K[:, 0, 0] = fx
399 | new_K[:, 1, 1] = fy
400 | new_K[:, 0, 2] = cx
401 | new_K[:, 1, 2] = cy
402 | return new_K
403 |
404 |
405 | def cropresize_backtransform_points2d(input_wh, boxes_2d_crop,
406 | output_wh, points_2d_in_output):
407 | bsz = input_wh.shape[0]
408 | assert output_wh.shape == (bsz, 2)
409 | assert input_wh.shape == (bsz, 2)
410 | assert points_2d_in_output.dim() == 3
411 |
412 | points_2d_normalized = points_2d_in_output / output_wh.unsqueeze(1)
413 | points_2d = boxes_2d_crop[:, [0, 1]].unsqueeze(1) + points_2d_normalized * input_wh.unsqueeze(1)
414 | return points_2d
415 |
--------------------------------------------------------------------------------
/lib/utils/urdf_robot.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
4 | import platform
5 | import numpy as np
6 | import pandas as pd
7 | import pyrender
8 | import torch
9 | from config import (BAXTER_DESCRIPTION_PATH, KUKA_DESCRIPTION_PATH,
10 | OWI_DESCRIPTION, OWI_KEYPOINTS_PATH,
11 | PANDA_DESCRIPTION_PATH, PANDA_DESCRIPTION_PATH_VISUAL)
12 | from dataset.const import JOINT_NAMES, LINK_NAMES
13 | from PIL import Image
14 | from utils.geometries import (quat_to_rotmat, rot6d_to_rotmat, rot9d_to_rotmat,
15 | rotmat_to_quat, rotmat_to_rot6d)
16 | from utils.mesh_renderer import RobotMeshRenderer, PandaArm
17 | from utils.urdfpytorch import URDF
18 |
19 | if platform.system() == "Linux":
20 | os.environ['PYOPENGL_PLATFORM'] = 'egl'
21 |
22 | class URDFRobot:
23 | def __init__(self,robot_type):
24 | self.robot_type = robot_type
25 | if self.robot_type == "panda":
26 | self.urdf_path = PANDA_DESCRIPTION_PATH
27 | self.urdf_path_visual = PANDA_DESCRIPTION_PATH_VISUAL
28 | self.dof = 8
29 | self.robot_for_render = PandaArm(self.urdf_path)
30 | elif self.robot_type == "kuka":
31 | self.urdf_path = KUKA_DESCRIPTION_PATH
32 | self.urdf_path_visual = KUKA_DESCRIPTION_PATH
33 | self.dof = 7
34 | self.robot_for_render = PandaArm(self.urdf_path)
35 | elif self.robot_type == "baxter":
36 | self.urdf_path = BAXTER_DESCRIPTION_PATH
37 | self.urdf_path_visual = BAXTER_DESCRIPTION_PATH
38 | self.dof = 15
39 | self.robot_for_render = PandaArm(self.urdf_path)
40 | elif self.robot_type == "owi":
41 | self.urdf_path = OWI_DESCRIPTION
42 | self.urdf_path_visual = OWI_DESCRIPTION
43 | self.dof = 4
44 | self.robot_for_render = None
45 | self.robot = URDF.load(self.urdf_path)
46 | self.robot_visual = URDF.load(self.urdf_path_visual)
47 | self.actuated_joint_names = JOINT_NAMES[self.robot_type]
48 | self.global_scale = 1.0
49 | self.device = None
50 | self.link_names, self.offsets = self.get_link_names_and_offsets()
51 |
52 | def get_link_names_and_offsets(self):
53 | if self.robot_type == "panda" or self.robot_type == "kuka":
54 | kp_offsets = torch.zeros((len(LINK_NAMES[self.robot_type]),3),dtype=torch.float).unsqueeze(0).unsqueeze(-1) * self.global_scale
55 | kp_offsets = kp_offsets.to(torch.float)
56 | return LINK_NAMES[self.robot_type], kp_offsets
57 | elif self.robot_type == "baxter":
58 | joint_name_to_joint = {joint.name: joint for joint in self.robot.joints}
59 | offsets = []
60 | link_names = []
61 | joint_names_for_links = [
62 | 'torso_t0', 'right_s0','left_s0', 'right_s1', 'left_s1',
63 | 'right_e0','left_e0', 'right_e1','left_e1','right_w0', 'left_w0',
64 | 'right_w1','left_w1','right_w2', 'left_w2','right_hand','left_hand'
65 | ]
66 | for joint_name in joint_names_for_links:
67 | joint = joint_name_to_joint[joint_name]
68 | offset = joint.origin[:3, -1]
69 | link_name = joint.parent
70 | link_names.append(link_name)
71 | offsets.append(offset)
72 | kp_offsets = torch.as_tensor(np.stack(offsets)).unsqueeze(0).unsqueeze(-1) * self.global_scale
73 | kp_offsets = kp_offsets.to(torch.float)
74 | return link_names, kp_offsets
75 | elif self.robot_type == "owi":
76 | keypoint_infos = pd.read_json(OWI_KEYPOINTS_PATH)
77 | kp_offsets = torch.as_tensor(np.stack(keypoint_infos['offset'])).unsqueeze(0).unsqueeze(-1).to(torch.float)
78 | return LINK_NAMES[self.robot_type], kp_offsets
79 | else:
80 | raise(NotImplementedError)
81 |
82 | def get_keypoints(self, jointcfgs, b2c_rot, b2c_trans):
83 |
84 | # jointcfgs, b2c_rot, b2c_trans all comes in batch (as model outputs)
85 | # b2c means base to camera
86 |
87 | batch_size = b2c_rot.shape[0]
88 | if b2c_rot.shape[1] == 6:
89 | rotmat = rot6d_to_rotmat(b2c_rot)
90 | elif b2c_rot.shape[1] == 4:
91 | rotmat = quat_to_rotmat(b2c_rot)
92 | elif b2c_rot.shape[1] == 9:
93 | rotmat = rot9d_to_rotmat(b2c_rot)
94 | else:
95 | raise NotImplementedError
96 | trans = b2c_trans.unsqueeze(dim=2)
97 | pad = torch.zeros((batch_size,1,4),dtype=torch.float).cuda()
98 | base2cam = torch.cat([rotmat,trans],dim=2).cuda()
99 | base2cam = torch.cat([base2cam,pad],dim=1).cuda()
100 | base2cam[:,3,3] = 1.0
101 | base2cam = base2cam.unsqueeze(1)
102 | TWL_base = self.get_TWL(jointcfgs).cuda()
103 | TWL = base2cam @ TWL_base
104 | pts = TWL[:, :, :3, :3] @ self.offsets.cuda() + TWL[:, :, :3, [-1]]
105 | return pts.squeeze(-1)
106 |
107 | def get_TWL(self, cfgs):
108 | fk = self.robot.link_fk_batch(cfgs, use_names=True)
109 | TWL = torch.stack([fk[link] for link in self.link_names]).permute(1, 0, 2, 3)
110 | TWL[..., :3, -1] *= self.global_scale
111 | return TWL
112 |
113 | def get_rotation_at_specific_root(self, jointcfgs, b2c_rot, b2c_trans, root = 0):
114 | if root == 0:
115 | return b2c_rot
116 | batch_size = b2c_rot.shape[0]
117 | if b2c_rot.shape[1] == 6:
118 | rotmat = rot6d_to_rotmat(b2c_rot)
119 | elif b2c_rot.shape[1] == 4:
120 | rotmat = quat_to_rotmat(b2c_rot)
121 | elif b2c_rot.shape[1] == 9:
122 | rotmat = rot9d_to_rotmat(b2c_rot)
123 | else:
124 | raise NotImplementedError
125 | trans = b2c_trans.unsqueeze(dim=2)
126 | pad = torch.zeros((batch_size,1,4),dtype=torch.float).cuda()
127 | base2cam = torch.cat([rotmat,trans],dim=2).cuda()
128 | base2cam = torch.cat([base2cam,pad],dim=1).cuda()
129 | base2cam[:,3,3] = 1.0
130 | base2cam = base2cam.unsqueeze(1)
131 | TWL_base = self.get_TWL(jointcfgs).cuda()
132 | TWL = base2cam @ TWL_base
133 | assert root < TWL.shape[1], (root, TWL.shape[1])
134 | if b2c_rot.shape[1] == 6:
135 | rotation = rotmat_to_rot6d(TWL[:, root, :3, :3]).cuda()
136 | elif b2c_rot.shape[1] == 4:
137 | rotation = rotmat_to_quat(TWL[:, root, :3, :3]).cuda()
138 | return rotation
139 |
140 |
141 | def get_keypoints_only_fk(self, jointcfgs):
142 |
143 | # only using joint angles to perform forward kinematics
144 | # the fk process is used when assuming the world frame is at the robot base, so rotation is identity and translation is origin/zeros
145 | # the output from this fk function is used for pnp process
146 |
147 | TWL = self.get_TWL(jointcfgs).cuda()
148 | pts = TWL[:, :, :3, :3] @ self.offsets.cuda() + TWL[:, :, :3, [-1]]
149 | return pts.squeeze(-1)
150 |
151 | def get_keypoints_only_fk_at_specific_root(self, jointcfgs, root=0):
152 |
153 | # only using joint angles to perform forward kinematics
154 | # the fk process is used when assuming the world frame is at the robot base, so rotation is identity and translation is origin/zeros
155 | # the output from this fk function is used for pnp process
156 |
157 | if root == 0:
158 | return self.get_keypoints_only_fk(jointcfgs)
159 | else:
160 | assert root > 0 and root < len(self.link_names)
161 |
162 | TWL_base = self.get_TWL(jointcfgs).cuda()
163 | TWL_root_inv = torch.linalg.inv(TWL_base[:,root:root+1,:,:])
164 | TWL = TWL_root_inv @ TWL_base
165 | pts = TWL[:, :, :3, :3] @ self.offsets.cuda() + TWL[:, :, :3, [-1]]
166 | return pts.squeeze(-1)
167 |
168 |
169 | def get_keypoints_root(self, jointcfgs, b2c_rot, b2c_trans, root = 0):
170 |
171 | # jointcfgs, b2c_rot, b2c_trans all comes in batch (as model outputs)
172 | # b2c here means *** root *** to camera
173 |
174 | if root == 0:
175 | return self.get_keypoints(jointcfgs, b2c_rot, b2c_trans)
176 | else:
177 | assert root > 0 and root < len(self.link_names)
178 |
179 | batch_size = b2c_rot.shape[0]
180 | if b2c_rot.shape[1] == 6:
181 | rotmat = rot6d_to_rotmat(b2c_rot)
182 | elif b2c_rot.shape[1] == 4:
183 | rotmat = quat_to_rotmat(b2c_rot)
184 | elif b2c_rot.shape[1] == 9:
185 | rotmat = rot9d_to_rotmat(b2c_rot)
186 | else:
187 | raise NotImplementedError
188 | trans = b2c_trans.unsqueeze(dim=2)
189 | pad = torch.zeros((batch_size,1,4),dtype=torch.float).cuda()
190 | base2cam = torch.cat([rotmat,trans],dim=2).cuda()
191 | base2cam = torch.cat([base2cam,pad],dim=1).cuda()
192 | base2cam[:,3,3] = 1.0
193 | base2cam = base2cam.unsqueeze(1)
194 | TWL_base = self.get_TWL(jointcfgs).cuda()
195 | TWL_root_inv = torch.linalg.inv(TWL_base[:,root:root+1,:,:])
196 | TWL_base = TWL_root_inv @ TWL_base
197 | TWL = base2cam @ TWL_base
198 | pts = TWL[:, :, :3, :3] @ self.offsets.cuda() + TWL[:, :, :3, [-1]]
199 | return pts.squeeze(-1)
200 |
201 | def set_robot_renderer(self, K_original, original_image_size=(480, 640), scale=0.5, device="cpu"):
202 |
203 | fx, fy, cx, cy = K_original[0,0]*scale, K_original[1,1]*scale, K_original[0,2]*scale, K_original[1,2]*scale
204 | image_size = (int(original_image_size[0]*scale), int(original_image_size[1]*scale))
205 |
206 | base_dir = os.path.dirname(self.urdf_path)
207 |
208 | mesh_files = [
209 | base_dir + "/meshes/visual/link0/link0.obj",
210 | base_dir + "/meshes/visual/link1/link1.obj",
211 | base_dir + "/meshes/visual/link2/link2.obj",
212 | base_dir + "/meshes/visual/link3/link3.obj",
213 | base_dir + "/meshes/visual/link4/link4.obj",
214 | base_dir + "/meshes/visual/link5/link5.obj",
215 | base_dir + "/meshes/visual/link6/link6.obj",
216 | base_dir + "/meshes/visual/link7/link7.obj",
217 | base_dir + "/meshes/visual/hand/hand.obj",
218 | ]
219 |
220 | focal_length = [-fx,-fy]
221 | principal_point = [cx, cy]
222 |
223 | robot_renderer = RobotMeshRenderer(
224 | focal_length=focal_length, principal_point=principal_point, image_size=image_size,
225 | robot=self.robot_for_render, mesh_files=mesh_files, device=device)
226 |
227 | return robot_renderer
228 |
229 | def get_robot_mesh_list(self, joint_angles, renderer):
230 |
231 | robot_meshes = []
232 | for joint_angle in joint_angles:
233 | if self.robot_type == "panda":
234 | joints = joint_angle[:-1].detach().cpu()
235 | else:
236 | joints = joint_angle.detach().cpu()
237 | robot_mesh = renderer.get_robot_mesh(joints).cuda()
238 | robot_meshes.append(robot_mesh)
239 |
240 | return robot_meshes
241 |
242 | def get_rendered_mask_single_image(self, rot, trans, robot_mesh, robot_renderer_gpu):
243 |
244 | R = rot6d_to_rotmat(rot).reshape(1,3,3)
245 | R = torch.transpose(R,1,2).cuda()
246 |
247 | T = trans.reshape(1,3).cuda()
248 |
249 | if T[0,-1] < 0:
250 | rendered_image = robot_renderer_gpu.silhouette_renderer(meshes_world=robot_mesh, R = -R, T = -T)
251 | else:
252 | rendered_image = robot_renderer_gpu.silhouette_renderer(meshes_world=robot_mesh, R = R, T = T)
253 |
254 | if torch.isnan(rendered_image).any():
255 | rendered_image = torch.nan_to_num(rendered_image)
256 |
257 | return rendered_image[..., 3]
258 |
259 | def get_rendered_mask_single_image_at_specific_root(self, joint_angles, rot, trans, robot_mesh, robot_renderer_gpu, root=0):
260 |
261 | if root == 0:
262 | return self.get_rendered_mask_single_image(rot, trans, robot_mesh, robot_renderer_gpu)
263 | else:
264 | rotmat = rot6d_to_rotmat(rot).cuda()
265 | trans = trans.unsqueeze(dim=1).cuda()
266 | pad = torch.zeros((1,4),dtype=torch.float).cuda()
267 | base2cam = torch.cat([rotmat,trans],dim=1).cuda()
268 | base2cam = torch.cat([base2cam,pad],dim=0).cuda()
269 | base2cam[3,3] = 1.0
270 | TWL_base = self.get_TWL(joint_angles.unsqueeze(0)).cuda().detach() # detach joint with rot/trans
271 | TWL_root_inv = torch.linalg.inv(TWL_base[:,root:root+1,:,:]).squeeze()
272 | new_base2cam = base2cam @ TWL_root_inv
273 | new_rot = rotmat_to_rot6d(new_base2cam[:3,:3])
274 | new_trans = new_base2cam[:3,3]
275 | return self.get_rendered_mask_single_image(new_rot, new_trans, robot_mesh, robot_renderer_gpu)
276 |
277 | def get_textured_rendering(self, joint, rot, trans, intrinsics=(320, 320, 320, 240), save_path=(None,None,None), original_image=None, root=0):
278 |
279 | if root != 0:
280 | rotmat = rot6d_to_rotmat(rot)
281 | trans = trans.unsqueeze(dim=1)
282 | pad = torch.zeros((1,4),dtype=torch.float)
283 | base2cam = torch.cat([rotmat,trans],dim=1)
284 | base2cam = torch.cat([base2cam,pad],dim=0)
285 | base2cam[3,3] = 1.0
286 | TWL_base = self.get_TWL(joint.unsqueeze(0))
287 | TWL_root_inv = torch.linalg.inv(TWL_base[:,root:root+1,:,:]).squeeze()
288 | new_base2cam = base2cam @ TWL_root_inv
289 | rot = rotmat_to_rot6d(new_base2cam[:3,:3])
290 | trans = new_base2cam[:3,3]
291 |
292 | save_path1, save_path2, save_path3 = save_path
293 | rotmat = rot6d_to_rotmat(rot)
294 | trans = trans.unsqueeze(dim=1)
295 | pad = torch.zeros((1,4),dtype=torch.float)
296 | camera_pose = torch.cat([rotmat,trans],dim=1)
297 | camera_pose = torch.cat([camera_pose,pad],dim=0)
298 | camera_pose[3,3] = 1.0
299 | joint = joint.numpy()
300 | camera_pose = camera_pose.numpy()
301 | rotation = np.array([[1,0,0,0],
302 | [0,-1,0,0],
303 | [0,0,-1,0],
304 | [0,0,0,1]])
305 | camera_pose = np.matmul(rotation,camera_pose)
306 | camera_pose = np.linalg.inv(camera_pose)
307 | fk = self.robot_visual.visual_trimesh_fk(cfg=joint)
308 | scene = pyrender.Scene()
309 | camera = pyrender.IntrinsicsCamera(*intrinsics)
310 | light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=3)
311 | light_pose = np.eye(4)
312 | light_pose[:3, 3] = np.array([0, -1, 1])
313 | scene.add(light, pose=light_pose)
314 | light_pose[:3, 3] = np.array([0, 1, 1])
315 | scene.add(light, pose=light_pose)
316 | light_pose[:3, 3] = np.array([1, 1, 2])
317 | scene.add(light, pose=light_pose)
318 | for tm in fk:
319 | pose = fk[tm]
320 | mesh = pyrender.Mesh.from_trimesh(tm, smooth=False)
321 | scene.add(mesh, pose=pose)
322 | scene.add(camera, pose=camera_pose)
323 | scene.add(light, pose=camera_pose)
324 | renderer = pyrender.OffscreenRenderer(viewport_width=640, viewport_height=480)
325 | color, depth = renderer.render(scene)
326 | rendered_img = Image.fromarray(np.uint8(color)).convert("RGBA")
327 | rendered_img.save(save_path1)
328 | original_img = Image.fromarray(np.transpose(np.uint8(original_image),(1,2,0))).convert("RGBA")
329 | original_img.save(save_path2)
330 | blend_ratio = 0.7
331 | blended_image = Image.blend(original_img, rendered_img, blend_ratio)
332 | blended_image.save(save_path3)
333 |
334 | def get_textured_rendering_individual(self, joint, rot, trans, intrinsics=(320, 320, 320, 240), root=0):
335 |
336 | if root != 0:
337 | rotmat = rot6d_to_rotmat(rot)
338 | trans = trans.unsqueeze(dim=1)
339 | pad = torch.zeros((1,4),dtype=torch.float)
340 | base2cam = torch.cat([rotmat,trans],dim=1)
341 | base2cam = torch.cat([base2cam,pad],dim=0)
342 | base2cam[3,3] = 1.0
343 | TWL_base = self.get_TWL(joint.unsqueeze(0))
344 | TWL_root_inv = torch.linalg.inv(TWL_base[:,root:root+1,:,:]).squeeze()
345 | new_base2cam = base2cam @ TWL_root_inv
346 | rot = rotmat_to_rot6d(new_base2cam[:3,:3])
347 | trans = new_base2cam[:3,3]
348 |
349 | rotmat = rot6d_to_rotmat(rot)
350 | trans = trans.unsqueeze(dim=1)
351 | pad = torch.zeros((1,4),dtype=torch.float)
352 | camera_pose = torch.cat([rotmat,trans],dim=1)
353 | camera_pose = torch.cat([camera_pose,pad],dim=0)
354 | camera_pose[3,3] = 1.0
355 | joint = joint.numpy()
356 | camera_pose = camera_pose.numpy()
357 | rotation = np.array([[1,0,0,0],
358 | [0,-1,0,0],
359 | [0,0,-1,0],
360 | [0,0,0,1]])
361 | camera_pose = np.matmul(rotation,camera_pose)
362 | camera_pose = np.linalg.inv(camera_pose)
363 | fk = self.robot_visual.visual_trimesh_fk(cfg=joint)
364 | scene = pyrender.Scene()
365 | camera = pyrender.IntrinsicsCamera(*intrinsics)
366 | # azure
367 | light = pyrender.PointLight(color=[1.5, 1.5, 1.5], intensity=2.6)
368 | # realsense, kinect
369 | # light = pyrender.PointLight(color=[1.4, 1.4, 1.4], intensity=2.4)
370 | # light_pose = np.eye(4)
371 | # light_pose[:3, 3] = np.array([0, 1, 0])
372 | # scene.add(light, pose=light_pose)
373 | # orb
374 | # light = pyrender.PointLight(color=[1.4, 1.4, 1.4], intensity=2.4)
375 | # light_pose = np.eye(4)
376 | # light_pose[:3, 3] = np.array([0, -1, 0])
377 | # scene.add(light, pose=light_pose)
378 |
379 | for tm in fk:
380 | pose = fk[tm]
381 | mesh = pyrender.Mesh.from_trimesh(tm, smooth=False)
382 | scene.add(mesh, pose=pose)
383 | scene.add(camera, pose=camera_pose)
384 | scene.add(light, pose=camera_pose)
385 | renderer = pyrender.OffscreenRenderer(viewport_width=640, viewport_height=480)
386 | color, depth = renderer.render(scene)
387 | rendered_img = Image.fromarray(np.uint8(color))
388 | return rendered_img
--------------------------------------------------------------------------------