Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import torch | |
| import matplotlib.pyplot as plt | |
| import matplotlib as mpl | |
| import os | |
| import cv2 | |
| from PIL import Image | |
| from .visualizer import Visualizer | |
| movable_imap = { | |
| 0: 'one_hand', | |
| 1: 'two_hands', | |
| 2: 'fixture', | |
| -100: 'n/a', | |
| } | |
| rigid_imap = { | |
| 1: 'yes', | |
| 0: 'no', | |
| 2: 'bad', | |
| -100: 'n/a', | |
| } | |
| kinematic_imap = { | |
| 0: 'freeform', | |
| 1: 'rotation', | |
| 2: 'translation', | |
| -100: 'n/a' | |
| } | |
| action_imap = { | |
| 0: 'free', | |
| 1: 'pull', | |
| 2: 'push', | |
| -100: 'n/a', | |
| } | |
| def draw_properties( | |
| output_path, | |
| movable, | |
| rigid, | |
| kinematic, | |
| action, | |
| ): | |
| fig = plt.figure() | |
| vis_phy = np.ones((80, 60, 3), dtype=np.uint8) * 255 | |
| plt.imshow(vis_phy) | |
| fig.axes[0].add_patch( | |
| mpl.patches.FancyBboxPatch( | |
| (10, 10), | |
| 40, | |
| 60, | |
| ec=np.array([79, 129, 189]) / 255.0, | |
| boxstyle=mpl.patches.BoxStyle("Round", pad=5), | |
| fill=False, | |
| linewidth=5, | |
| ) | |
| ) | |
| fontsize = 20 | |
| if rigid_imap[rigid] == 'yes': | |
| rigid_text = 'Rigid: Yes' | |
| else: | |
| rigid_text = 'Rigid: No' | |
| plt.text(30, 20, rigid_text, fontsize=fontsize, horizontalalignment='center') | |
| if movable_imap[movable] == 'one_hand': | |
| movable_text = "Move: 1 hand" | |
| elif movable_imap[movable] == 'two_hands': | |
| movable_text = "Move: 2 hands" | |
| else: | |
| movable_text = 'Move: Fixture' | |
| plt.text(30, 35, movable_text, fontsize=fontsize, horizontalalignment='center') | |
| if kinematic_imap[kinematic] == 'rotation': | |
| kinematic_text = 'Motion: Rot' | |
| elif kinematic_imap[kinematic] == 'translation': | |
| kinematic_text = 'Motion: Trans' | |
| else: | |
| kinematic_text = 'Motion: Free' | |
| plt.text(30, 50, kinematic_text, fontsize=fontsize, horizontalalignment='center') | |
| if action_imap[action] == 'pull': | |
| action_text = 'Action: Pull' | |
| elif action_imap[action] == 'push': | |
| action_text = 'Action: Push' | |
| else: | |
| action_text = 'Action: Free' | |
| plt.text(30, 65, action_text, fontsize=fontsize, horizontalalignment='center') | |
| plt.axis('off') | |
| plt.savefig(output_path, bbox_inches='tight', pad_inches=0) | |
| plt.close(fig) | |
| def draw_localization(rgb, output_path, bbox, mask, axis, colors=None, alpha=0.5): | |
| if colors is None: | |
| # colors = np.array([[246, 95, 16], ]) / 255.0 # orange | |
| #colors = np.array([[227, 13, 255], ]) / 255.0 # purple | |
| colors = np.array([[178, 13, 255], ]) / 255.0 # purple | |
| pred_entry = { | |
| 'keypoint': None, | |
| 'bbox': bbox, | |
| 'mask': mask, | |
| 'affordance': None, | |
| 'move': None, | |
| 'rigid': None, | |
| 'kinematic': None, | |
| 'pull_or_push': None, | |
| 'axis': axis, | |
| } | |
| instances = [pred_entry] | |
| vis = Visualizer(rgb) | |
| vis.overlay_instances(instances, assigned_colors=colors, alpha=alpha) | |
| vis.output.save(output_path) | |
| def draw_affordance(rgb, output_path, affordance, alpha=0.5): | |
| # normalize affordance | |
| aff_min = affordance.min() | |
| aff_max = affordance.max() | |
| affordance = (affordance - aff_min) / (aff_max - aff_min) | |
| # convert rgb to gray scale for better visualization | |
| rgb_gray = rgb.copy() | |
| # rgb_gray[:, :, 0] = rgb.mean(axis=2) | |
| # rgb_gray[:, :, 1] = rgb.mean(axis=2) | |
| # rgb_gray[:, :, 2] = rgb.mean(axis=2) | |
| heatmap_img = cv2.applyColorMap((affordance * 255.0).astype(np.uint8), cv2.COLORMAP_HOT)[:, :, ::-1] | |
| vis = cv2.addWeighted(heatmap_img, alpha, rgb_gray, 1-alpha, 0) | |
| return Image.fromarray(vis) #.save(output_path) |