| 61 | return data |
| 62 | |
| 63 | def set_startpoint_to_zero(data): |
| 64 | num_obj = data.shape[-1] // 13 |
| 65 | end_loc = np.zeros([data.shape[0], data.shape[1], num_obj, 3]) |
| 66 | end_loc[..., :18, :2] += np.copy(data[:, :1, None, :2]) |
| 67 | end_loc[..., 19:, :2] += np.copy(data[:, :1, None, :2]) |
| 68 | end_loc = end_loc.reshape([data.shape[0], data.shape[1], num_obj * 3]) |
| 69 | data[..., :num_obj*3] = np.copy(data[..., :num_obj*3]) - end_loc |
| 70 | return data |
| 71 | |
| 72 | def diffusion_planner(text, pre_seq, waypoint, motion_length): |
| 73 | if waypoint is not None and pre_seq is not None: |