File size: 5,236 Bytes
987ed1b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | from typing import Dict, Sequence, Union, Optional
import numpy as np
import skimage.transform as st
import pymunk
import pygame
from matplotlib import cm
import cv2
from diffusion_policy.env.pusht.pymunk_override import DrawOptions
def farthest_point_sampling(points: np.ndarray, n_points: int, init_idx: int):
"""
Naive O(N^2)
"""
assert(n_points >= 1)
chosen_points = [points[init_idx]]
for _ in range(n_points-1):
cpoints = np.array(chosen_points)
all_dists = np.linalg.norm(points[:,None,:] - cpoints[None,:,:], axis=-1)
min_dists = all_dists.min(axis=1)
next_idx = np.argmax(min_dists)
next_pt = points[next_idx]
chosen_points.append(next_pt)
result = np.array(chosen_points)
return result
class PymunkKeypointManager:
def __init__(self,
local_keypoint_map: Dict[str, np.ndarray],
color_map: Optional[Dict[str, np.ndarray]]=None):
"""
local_keypoint_map:
"<attribute_name>": (N,2) floats in object local coordinate
"""
if color_map is None:
cmap = cm.get_cmap('tab10')
color_map = dict()
for i, key in enumerate(local_keypoint_map.keys()):
color_map[key] = (np.array(cmap.colors[i]) * 255).astype(np.uint8)
self.local_keypoint_map = local_keypoint_map
self.color_map = color_map
@property
def kwargs(self):
return {
'local_keypoint_map': self.local_keypoint_map,
'color_map': self.color_map
}
@classmethod
def create_from_pusht_env(cls, env, n_block_kps=9, n_agent_kps=3, seed=0, **kwargs):
rng = np.random.default_rng(seed=seed)
local_keypoint_map = dict()
for name in ['block','agent']:
self = env
self.space = pymunk.Space()
if name == 'agent':
self.agent = obj = self.add_circle((256, 400), 15)
n_kps = n_agent_kps
else:
self.block = obj = self.add_tee((256, 300), 0)
n_kps = n_block_kps
self.screen = pygame.Surface((512,512))
self.screen.fill(pygame.Color("white"))
draw_options = DrawOptions(self.screen)
self.space.debug_draw(draw_options)
# pygame.display.flip()
img = np.uint8(pygame.surfarray.array3d(self.screen).transpose(1, 0, 2))
obj_mask = (img != np.array([255,255,255],dtype=np.uint8)).any(axis=-1)
tf_img_obj = cls.get_tf_img_obj(obj)
xy_img = np.moveaxis(np.array(np.indices((512,512))), 0, -1)[:,:,::-1]
local_coord_img = tf_img_obj.inverse(xy_img.reshape(-1,2)).reshape(xy_img.shape)
obj_local_coords = local_coord_img[obj_mask]
# furthest point sampling
init_idx = rng.choice(len(obj_local_coords))
obj_local_kps = farthest_point_sampling(obj_local_coords, n_kps, init_idx)
small_shift = rng.uniform(0, 1, size=obj_local_kps.shape)
obj_local_kps += small_shift
local_keypoint_map[name] = obj_local_kps
return cls(local_keypoint_map=local_keypoint_map, **kwargs)
@staticmethod
def get_tf_img(pose: Sequence):
pos = pose[:2]
rot = pose[2]
tf_img_obj = st.AffineTransform(
translation=pos, rotation=rot)
return tf_img_obj
@classmethod
def get_tf_img_obj(cls, obj: pymunk.Body):
pose = tuple(obj.position) + (obj.angle,)
return cls.get_tf_img(pose)
def get_keypoints_global(self,
pose_map: Dict[set, Union[Sequence, pymunk.Body]],
is_obj=False):
kp_map = dict()
for key, value in pose_map.items():
if is_obj:
tf_img_obj = self.get_tf_img_obj(value)
else:
tf_img_obj = self.get_tf_img(value)
kp_local = self.local_keypoint_map[key]
kp_global = tf_img_obj(kp_local)
kp_map[key] = kp_global
return kp_map
def draw_keypoints(self, img, kps_map, radius=1):
scale = np.array(img.shape[:2]) / np.array([512,512])
for key, value in kps_map.items():
color = self.color_map[key].tolist()
coords = (value * scale).astype(np.int32)
for coord in coords:
cv2.circle(img, coord, radius=radius, color=color, thickness=-1)
return img
def draw_keypoints_pose(self, img, pose_map, is_obj=False, **kwargs):
kp_map = self.get_keypoints_global(pose_map, is_obj=is_obj)
return self.draw_keypoints(img, kps_map=kp_map, **kwargs)
def test():
from diffusion_policy.environment.push_t_env import PushTEnv
from matplotlib import pyplot as plt
env = PushTEnv(headless=True, obs_state=False, draw_action=False)
kp_manager = PymunkKeypointManager.create_from_pusht_env(env=env)
env.reset()
obj_map = {
'block': env.block,
'agent': env.agent
}
obs = env.render()
img = obs.astype(np.uint8)
kp_manager.draw_keypoints_pose(img=img, pose_map=obj_map, is_obj=True)
plt.imshow(img)
|