def visualize_example(example): img = example["image"].numpy() cms = example["predicted_confidence_maps"].numpy() pts_gt = example["instances"].numpy()[0] pts_pr = example["predicted_points"].numpy() scale = 1.0 if img.shape[0] < 512: scale = 2.0 if img.shape[0] < 256: scale = 4.0 fig = plot_img(img, dpi=72 * scale, scale=scale) plot_confmaps(cms, output_scale=cms.shape[0] / img.shape[0]) plot_peaks(pts_gt, pts_pr, paired=True) return fig
def visualize_pafs_example(example): img = example["image"].numpy() pafs = example["predicted_part_affinity_fields"].numpy() scale = 1.0 if img.shape[0] < 512: scale = 2.0 if img.shape[0] < 256: scale = 4.0 fig = plot_img(img, dpi=72 * scale, scale=scale) plot_pafs( pafs, output_scale=pafs.shape[0] / img.shape[0], stride=1, scale=8.0, width=1.0, ) return fig