-
Notifications
You must be signed in to change notification settings - Fork 0
/
param_selection_demo.py
157 lines (131 loc) · 6.9 KB
/
param_selection_demo.py
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
148
149
150
151
152
153
154
155
156
157
from environments.realrobot_env_multi import RealRobotEnvMulti
import tkinter as tk
from PIL import Image, ImageTk
import numpy as np
import cv2
from utils.detection_utils_eeg import DetectionUtils
class ParamSelectionDemo:
def __init__(self):
self.pix_pos = None
self.chose_xy = False
self.chose_z = False
self.canvas = None
self.window = None
self.env = RealRobotEnvMulti()
self.detection_utils = DetectionUtils()
def demo_pick(self):
# brings up a UI to click positions on images to select parameters for "pick" skill
# take images and display cam 2 image
dummy_action = np.zeros(self.env.skill.num_skills)
imgs = self.env.get_image_observations(action=dummy_action, img_as_list=False, save_images=False)
param_img2 = imgs["param_image2"]
param_img0 = imgs["param_image0"]
pil_image0 = Image.fromarray(param_img0[:,:,::-1])
pil_image2 = Image.fromarray(param_img2[:,:,::-1])
############### choose x and y ################
# Setup Tkinter window and canvas with topdown view (cam 2 for param selection) - NOTE this block can be replaced with
self.window = tk.Tk()
self.canvas = tk.Canvas(self.window, width=pil_image2.width, height=pil_image2.height)
self.canvas.pack()
photo = ImageTk.PhotoImage(pil_image2)
self.canvas.create_image(0, 0, image=photo, anchor=tk.NW)
self.canvas.bind("<Button 1>", self.get_pixel_position) # bind mouseclick event to canvas
self.window.mainloop() # start Tkinter event loop
# get xy world position from pixel position
world_xy = self.detection_utils.get_world_xy_from_topdown_view(pix_coords=self.pix_pos, img_array=param_img2)
############### choose z ###############
# discretize z coordinate and save image visualizing the discrete points
pix_pts, world_pts = self.detection_utils.get_points_on_z(world_xy=world_xy, img_array=param_img0, camera_id=0, max_height=0.3)
projection_img = cv2.imread("projections.png")
projection_pil_img = Image.fromarray(projection_img[:,:,::-1])
# Setup Tkinter window and canvas with side view (cam 0 for param selection)
self.window = tk.Tk()
self.canvas = tk.Canvas(self.window, width=projection_pil_img.width, height=projection_pil_img.height)
self.canvas.pack()
photo = ImageTk.PhotoImage(projection_pil_img)
self.canvas.create_image(0, 0, image=photo, anchor=tk.NW)
self.canvas.bind("<Button 1>", self.get_pixel_position)
self.window.mainloop()
# find the point in the discrete pixel points closest to clicked point
selected_pix_pos = np.array([self.pix_pos[1], self.pix_pos[0]])
idx = np.argmin(np.linalg.norm(pix_pts - selected_pix_pos, axis=1))
world_z = world_pts[idx][2]
# execute pick
skill_selection_vec = np.zeros(self.env.num_skills)
skill_selection_vec[0] = 1
params = [world_xy[0], world_xy[1], world_z]
print("selected world coordinates", params)
action = np.concatenate([skill_selection_vec, params])
obs, reward, done, info = self.env.step(action)
def get_pixel_position(self, event):
# Get the position of the mouse click in the image
x, y = event.x, event.y
# Print the pixel position
print("Pixel position:", x, y)
self.pix_pos = x,y
self.window.destroy()
class ParamSelectionDemoStandalone:
def __init__(self):
self.pix_pos = None
self.chose_xy = False
self.chose_z = False
self.canvas = None
self.window = None
self.detection_utils = DetectionUtils()
def demo_pick(self, topdown_img_path, sideview_img_path, side_camera_id):
# brings up a UI to click positions on images to select parameters for "pick" skill
# reads images
topdown_img = cv2.imread(topdown_img_path)
sideview_img = cv2.imread(sideview_img_path)
topdown_pil_image = Image.fromarray(topdown_img[:,:,::-1])
############### choose x and y ################
# Setup Tkinter window and canvas with topdown view (cam 2 for param selection) - NOTE this block can be replaced with cursor control
self.window = tk.Tk()
self.canvas = tk.Canvas(self.window, width=topdown_pil_image.width, height=topdown_pil_image.height)
self.canvas.pack()
photo = ImageTk.PhotoImage(topdown_pil_image)
self.canvas.create_image(0, 0, image=photo, anchor=tk.NW)
self.canvas.bind("<Button 1>", self.get_pixel_position) # bind mouseclick event to canvas
self.window.mainloop() # start Tkinter event loop
# get xy world position from pixel position
world_xy = self.detection_utils.get_world_xy_from_topdown_view(pix_coords=self.pix_pos, img_array=topdown_img)
############### choose z ###############
# discretize z coordinate and save image visualizing the discrete points
pix_pts, world_pts = self.detection_utils.get_points_on_z(world_xy=world_xy, img_array=sideview_img, camera_id=side_camera_id, max_height=0.3)
projection_img = cv2.imread("projections.png")
projection_pil_img = Image.fromarray(projection_img[:,:,::-1])
# Setup Tkinter window and canvas with side view (cam 0 for param selection)
self.window = tk.Tk()
self.canvas = tk.Canvas(self.window, width=projection_pil_img.width, height=projection_pil_img.height)
self.canvas.pack()
photo = ImageTk.PhotoImage(projection_pil_img)
self.canvas.create_image(0, 0, image=photo, anchor=tk.NW)
self.canvas.bind("<Button 1>", self.get_pixel_position)
self.window.mainloop()
# find the point in the discrete pixel points closest to clicked point
selected_pix_pos = np.array([self.pix_pos[1], self.pix_pos[0]])
idx = np.argmin(np.linalg.norm(pix_pts - selected_pix_pos, axis=1))
world_z = world_pts[idx][2]
# execute pick
env_num_skills = 12
skill_selection_vec = np.zeros(env_num_skills)
skill_selection_vec[0] = 1
params = [world_xy[0], world_xy[1], world_z]
print("selected world coordinates", params)
return params
# action = np.concatenate([skill_selection_vec, params])
# obs, reward, done, info = self.env.step(action)
def get_pixel_position(self, event):
# Get the position of the mouse click in the image
x, y = event.x, event.y
# Print the pixel position
print("Pixel position:", x, y)
self.pix_pos = x,y
self.window.destroy()
def main():
demo = ParamSelectionDemoStandalone()
demo.demo_pick(topdown_img_path="param_selection_img2.png", sideview_img_path="param_selection_img0.png", side_camera_id=0)
# demo = ParamSelectionDemo()
# demo.demo_pick()
if __name__ == "__main__":
main()