| 4 | |
| 5 | |
| 6 | def transform2origin(position_tensor, scale = 1.0): |
| 7 | min_pos = torch.min(position_tensor, 0)[0] |
| 8 | max_pos = torch.max(position_tensor, 0)[0] |
| 9 | max_diff = torch.max(max_pos - min_pos) |
| 10 | original_mean_pos = (min_pos + max_pos) / 2.0 |
| 11 | scale = scale / max_diff |
| 12 | original_mean_pos = original_mean_pos.to(device="cuda") |
| 13 | scale = scale.to(device="cuda") |
| 14 | new_position_tensor = (position_tensor - original_mean_pos) * scale |
| 15 | |
| 16 | return new_position_tensor, scale, original_mean_pos |
| 17 | |
| 18 | |
| 19 | def undotransform2origin(position_tensor, scale, original_mean_pos): |