yfzhoucs/TinyLanguageRobots
1
1import gradio as gr2import torch3from tiny_ur5 import TinyUR5Env4import yaml5from initializer import Initializer6import random7import string8import imageio9from skimage import img_as_ubyte10from test_model import model_forward_fn11from PIL import Image12 13 14def load_model(ckpt, method, device):15 if method == 'bcz':16 from models.film_model import Backbone17 # model = Backbone(img_size=224, num_traces_out=4, embedding_size=256, num_weight_points=10, input_nc=3, device=device)18 model = Backbone(img_size=224, num_traces_out=8, embedding_size=256, num_weight_points=12, input_nc=3, device=device)19 model.load_state_dict(torch.load(ckpt, map_location=device)['model'], strict=True)20 # model = model.cpu()21 model = model.to(device)22 return model23 elif method == 'ours':24 # import tinyur5.models.backbone_rgbd_sub_attn_tinyur5.Backbone as Backbone25 # import tinyur526 # import models.backbone_rgbd_sub_attn_tinyur5.Backbone27 from models.backbone_rgbd_sub_attn_tinyur5 import Backbone28 # from tinyur5.models.backbone_rgbd_sub_attn_tinyur5 import Backbone29 model = Backbone(img_size=224, embedding_size=256, num_traces_out=2, num_joints=8, num_weight_points=12, input_nc=3, device=device)30 model.load_state_dict(torch.load(ckpt, map_location=device)['model'], strict=True)31 model = model.to(device)32 return model33 34device = torch.device('cpu')35# ckpt = '580000.pth'36ckpt = '160000.pth'37print('start loading model')38model = load_model(ckpt, 'ours', device)39print('model loaded')40 41 42 43with gr.Blocks() as demo:44 45 46 state = gr.State()47 48 # with open('config.yaml', "r") as stream:49 # try:50 # config = yaml.safe_load(stream)51 # # print(config, type(config))52 # except yaml.YAMLError as exc:53 # print(exc)54 55 # initializer = Initializer(config)56 57 # config, task = initializer.get_config_and_task()58 # sentence = initializer.get_sentence()59 # env = TinyUR5Env(config)60 61 62 def init(environment):63 if environment == 'original':64 config_file = 'config.yaml'65 else:66 config_file = 'config_stable_diffusion.yaml'67 # with open('config.yaml', "r") as stream:68 with open(config_file, "r") as stream:69 try:70 config = yaml.safe_load(stream)71 # print(config, type(config))72 except yaml.YAMLError as exc:73 print(exc)74 75 if environment == 'original':76 initializer = Initializer(config, obj_num_low=3, obj_num_high=5)77 else:78 initializer = Initializer(config, obj_num_low=1, obj_num_high=2)79 80 config, task = initializer.get_config_and_task()81 sentence = initializer.get_sentence()82 env = TinyUR5Env(config)83 init_img = env.render('rgb_array')84 current_state = {85 'env': env,86 'id': ''.join(random.choice(string.ascii_lowercase + string.ascii_uppercase + string.ascii_letters) for i in range(20))87 }88 return init_img, current_state89 90 91 def exec(sentence, current_state, resolution):92 env = current_state['env']93 img = env.render('rgb_array')94 95 imgs = []96 time_step = 097 while time_step < 150:98 actions = model_forward_fn(env, model, sentence, 'ours', device)99 # for i in range(actions.shape[-1]):100 for i in range(15, 50):101 action = actions[:, i]102 observation, reward, done, info = env.step(action, eef_z=80)103 img = env.render('rgb_array')104 img = Image.fromarray(img)105 if resolution == 'low(3 sec)':106 img = img.resize((240, 140))107 elif resolution == 'mid(5 sec)':108 img = img.resize((480, 280))109 elif resolution == 'high(7 sec)':110 img = img.resize((720, 420))111 # imgs.append(Image.fromarray(img))112 if time_step % 12 == 0:113 imgs.append(img)114 time_step += 1115 print(time_step)116 env.close()117 118 # context = {}119 # is_success, buffer = cv2.imencode(".jpg", cv2.cvtColor(img, cv2.COLOR_RGB2BGR))120 # img_buffer = BytesIO()121 # imgs[0].save(img_buffer, save_all=True, append_images=imgs[1:], duration=100, loop=0)122 # img = base64.b64encode(img_buffer.getvalue()).decode('utf-8')123 124 # imageio.mimsave(os.path.join('tinyur5/static/', request.session['id'] + '.gif') , [img_as_ubyte(frame) for frame in imgs], 'GIF', fps=20)125 # with open(os.path.join('tinyur5/static/', request.session['id'] + '.gif'), "rb") as gif_file:126 # img = format(base64.b64encode(gif_file.read()).decode())127 img_id = ''.join(random.choice(string.ascii_lowercase + string.ascii_uppercase + string.ascii_letters) for i in range(20))128 imageio.mimsave(img_id+'.gif', [img_as_ubyte(frame) for frame in imgs], 'GIF', fps=10)129 130 131 132 133 134 img = img_id+'.gif'135 next_state = {136 'id': current_state['id'],137 'env': env138 }139 return env.render('rgb_array'), img, next_state140 141 142 with gr.Row():143 with gr.Column(scale=4):144 instruction = gr.Text(label="""Input an Instruction Here:""", placeholder='Push XXX to the right / Rotate XXX')145 with gr.Column(scale=2):146 resolution = gr.Radio(147 label='Image Quality', 148 choices=['low(3 sec)', 'mid(5 sec)', 'high(7 sec)'],149 value='low(3 sec)')150 with gr.Column(scale=1):151 environment = gr.Radio(152 label='Environment', 153 choices=['original', 'stable diffusion'],154 value='original')155 with gr.Row():156 action = gr.Button(value='Action!')157 with gr.Row():158 init_img_placeholder = gr.Image()159 gif_img_placeholder = gr.Image()160 161 with gr.Row():162 load_env = gr.Button(value='Reload Simulator')163 with gr.Row():164 with gr.Column():165 illustration = gr.Markdown(166 # label='Try Commanding the Robot Yourself!', 167 value=168 """169 ## Commanding the Robot Yourself!170 (1) Type in some instructions in the instruction box at the top. 171 (2) Hit 'Action!' button to start executing your instruction. 172 (3) Hit 'Reload Simulator' button if you want to re-initialize the simulator.173 ## Try the images generated from stable diffusion!174 Click on the 'stable diffusion' radio for initializing the environment by images generated by stable diffusion.175 """,176 # lines=3,177 # interactive=False178 )179 with gr.Column():180 illustration = gr.Markdown(181 # label='Sample instructions:', 182 value=183 """184 ## Sample Instructions: 185 The robot can support pushing the objects in 4 directions, as well as rotating them: 186 ```187 \u2022 Push the apple to the right 188 \u2022 Rotate the watermelon clockwise 189 \u2022 Move the clock backwards 190 ```191 """,192 # lines=4,193 # interactive=False194 )195 196 load_env.click(197 init,198 inputs=[environment],199 outputs=[init_img_placeholder, state],200 show_progress=True201 )202 203 action.click(204 exec,205 inputs=[instruction, state, resolution],206 outputs=[init_img_placeholder, gif_img_placeholder, state],207 show_progress=True208 )209 demo.load(210 init,211 inputs=[environment],212 outputs=[init_img_placeholder, state],213 show_progress=True)214 215 environment.change(216 init,217 inputs=[environment],218 outputs=[init_img_placeholder, state],219 show_progress=True220 )221 222 223 224demo.launch(share=False)225