| 13 | |
| 14 | |
| 15 | def load_data(data_dir: str, nominal_freq: int): |
| 16 | with open(data_dir + "/data.pkl", "rb") as f: |
| 17 | data = pickle.load(f) |
| 18 | ee_pose = data["ee_pose"] |
| 19 | force_data = data["force"] |
| 20 | |
| 21 | ee_timestamps = ee_pose[:, 0] |
| 22 | duration = ee_timestamps[-1] - ee_timestamps[0] |
| 23 | num_frames = int(duration * nominal_freq) |
| 24 | start_timestamps = [ee_timestamps[0], force_data[0, 0]] |
| 25 | end_timestamps = [ee_timestamps[-1], force_data[-1, 0]] |
| 26 | |
| 27 | start_timestamp = max(start_timestamps) |
| 28 | end_timestamp = min(end_timestamps) |
| 29 | |
| 30 | interpolating_timestamps = np.linspace(start_timestamp, end_timestamp, num_frames) |
| 31 | print(f"start_timestamp: {start_timestamp}, end_timestamp: {end_timestamp}") |
| 32 | print(f"duration: {duration}, num_frames: {num_frames}") |
| 33 | |
| 34 | ee_pose = interp1d(ee_timestamps, ee_pose[:, 1:4], axis=0)(interpolating_timestamps) |
| 35 | forces = interp1d(force_data[:, 0], force_data[:, 1:4], axis=0)( |
| 36 | interpolating_timestamps |
| 37 | ) |
| 38 | # Invert normal force to get positive forces |
| 39 | forces[:, -1] = forces[:, -1] * -1 |
| 40 | return ee_pose, forces |
| 41 | |
| 42 | |
| 43 | def extract_strokes(in_contact: np.ndarray): |