| 19 | |
| 20 | |
| 21 | def plot_correlation(forces_gt, forces_pred): |
| 22 | corr_metric = [] |
| 23 | for i in range(3): |
| 24 | corr = stats.pearsonr(forces_gt[:, i], forces_pred[:, i]) |
| 25 | corr_metric.append(corr[0]) |
| 26 | |
| 27 | correlation_fig = plt.figure(figsize=(20, 5)) |
| 28 | axs: np.ndarray = correlation_fig.subplots(1, 3) |
| 29 | for i, (force_gt, force_pred) in enumerate(zip(forces_gt.T, forces_pred.T)): |
| 30 | axs[i].scatter( |
| 31 | force_gt, |
| 32 | force_pred, |
| 33 | s=2, |
| 34 | color=colors[i], |
| 35 | label=f"r={corr_metric[i]:.3f}", |
| 36 | ) |
| 37 | axs[i].set_xlabel("Ground Truth (N)") |
| 38 | axs[i].set_ylabel("Prediction (N)") |
| 39 | axs[i].set_title(f"Force {['X', 'Y', 'Z'][i]}") |
| 40 | axs[i].grid(True) |
| 41 | # plot 1:1 line |
| 42 | axs[i].plot( |
| 43 | [force_gt.min(), force_gt.max()], |
| 44 | [force_gt.min(), force_gt.max()], |
| 45 | "--", |
| 46 | color="gray", |
| 47 | ) |
| 48 | axs[i].legend() |
| 49 | # return correlation_fig, axs |
| 50 | |
| 51 | img_buf = io.BytesIO() |
| 52 | plt.savefig(img_buf, format="png") |
| 53 | plt.close("all") |
| 54 | im = Image.open(img_buf) |
| 55 | return im |
| 56 | |
| 57 | |
| 58 | def plot_forces_error(forces_gt, forces_pred, n_bins=100, n_std=3): |