import gradio as gr from models import build_model from PIL import Image import numpy as np import torchvision import math import ninja import torch from tqdm import trange import imageio import requests import argparse import imageio import spaces from scipy.spatial.transform import Rotation from gradio_draggable import Draggable checkpoint = 'clevr.pth' state = torch.load(checkpoint, map_location='cpu') G = build_model(**state['model_kwargs_init']['generator_smooth']) o0, o1 = G.load_state_dict(state['models']['generator_smooth'], strict=False) G.eval().cuda() G.backbone.synthesis.input.x_offset =0 G.backbone.synthesis.input.y_offset =0 G_kwargs= dict(noise_mode='const', fused_modulate=False, impl='cuda', fp16_res=None) print('prepare finish', flush=True) COLOR_NAME_LIST = ['cyan', 'green', 'purple', 'red', 'yellow', 'gray', 'purple', 'blue'] SHAPE_NAME_LIST = ['cube', 'sphere', 'cylinder'] MATERIAL_NAME_LIST = ['rubber', 'metal'] canvas_x = 800 canvas_y = 200 batch_size = 1 code = torch.randn(1, G.z_dim).cuda() to_pil = torchvision.transforms.ToPILImage() RT = torch.tensor([[ -1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.5000, -0.8660, 10.3923, 0.0000, -0.8660, -0.5000, 6.0000, 0.0000, 0.0000, 0.0000, 1.0000, 262.5000, 0.0000, 32.0000, 0.0000, 262.5000, 32.0000, 0.0000, 0.0000, 1.0000]], device='cuda') obj_dict = {} # init fake_bevs = torch.zeros([1, 14, 256, 256], device='cuda').float() _ = G(code, RT, fake_bevs) def trans(x, y, z, length): w = h = length x = 0.5 * w - 128 + 256 - (x/9 + .5) * 256 y = 0.5 * h - 128 + (y/9 + .5) * 256 z = z / 9 * 256 return x, y, z def objs_to_canvas(lst, length=256, scale = 2.6): objs = [] for each in lst: x, y, obj_id = each['x'], each['y'], each['id'] if obj_id not in obj_dict: color = np.random.choice(COLOR_NAME_LIST) shape = 'cube' material = 'rubber' rot = 0 obj_dict[obj_id] = [color, shape, material, rot] color, shape, material, rot = obj_dict[obj_id] x = -x / canvas_x * 16 y = y / canvas_y * 2 y *= 2 x += 1.0 y -= 1.5 z = 0.35 objs.append([x, y, z, shape, color, material, rot]) h, w = length, int(length *scale) nc = 14 canvas = np.zeros([h, w, nc]) xx = np.ones([h,w]).cumsum(0) yy = np.ones([h,w]).cumsum(1) for x, y, z, shape, color, material, rot in objs: y, x, z = trans(x, y, z, length) feat = [0] * nc feat[0] = 1 feat[COLOR_NAME_LIST.index(color) + 1] = 1 feat[SHAPE_NAME_LIST.index(shape) + 1 + len(COLOR_NAME_LIST)] = 1 feat[MATERIAL_NAME_LIST.index(material) + 1 + len(COLOR_NAME_LIST) + len(SHAPE_NAME_LIST)] = 1 feat = np.array(feat) rot_sin = np.sin(rot / 180 * np.pi) rot_cos = np.cos(rot / 180 * np.pi) if shape == 'cube': mask = (np.abs(+rot_cos * (xx-x) + rot_sin * (yy-y)) <= z) * \ (np.abs(-rot_sin * (xx-x) + rot_cos * (yy-y)) <= z) else: mask = ((xx-x)**2 + (y-yy)**2) ** 0.5 <= z canvas[mask] = feat canvas = np.transpose(canvas, [2, 0, 1]).astype(np.float32) return canvas @spaces.GPU @torch.no_grad() def predict_local_view(lst): canvas = torch.tensor(objs_to_canvas(lst)).cuda()[None] bevs = canvas[..., 0: 0+256] print(code.shape, RT.shape, bevs.shape) gen = G(code, RT, bevs) rgb = gen['gen_output']['image'][0] * .5 + .5 return to_pil(rgb) @spaces.GPU @torch.no_grad() def predict_local_view_video(lst): canvas = torch.tensor(objs_to_canvas(lst)).cuda()[None] bevs = canvas[..., 0: 0+256] RT_array = np.array(RT[0].cpu()) rot = RT_array[:16].reshape(4,4) trans = RT_array[16:] rot_new = rot.copy() r = Rotation.from_matrix(rot[:3, :3]) angles = r.as_euler("zyx",degrees=True) v_mean, h_mean = angles[1], angles[2] writer = imageio.get_writer('tmp.mp4', fps=25) for t in np.linspace(0, 1, 50): angles[1] = 0.5 * np.cos(t * 2 * math.pi) + v_mean angles[2] = 1 * np.sin(t * 2 * math.pi) + h_mean r = Rotation.from_euler("zyx",angles,degrees=True) rot_new[:3,:3] = r.as_matrix() new_RT = torch.tensor(np.concatenate([rot_new.flatten(), trans])[None]).cuda().float() gen = G(code, new_RT, bevs) rgb = gen['gen_output']['image'][0] * .5 + .5 writer.append_data(np.array(to_pil(rgb))) writer.close() return 'tmp.mp4' @spaces.GPU @torch.no_grad() def predict_global_view(lst): canvas = torch.tensor(objs_to_canvas(lst)).cuda()[None] length = canvas.shape[-1] lines = [] for i in trange(0, length - 256, 10): bevs = canvas[..., i: i+256] gen = G(code, RT, bevs) start = 128 if i > 0 else 0 lines.append(gen['gen_output']['image'][0, ..., start:128+32]) rgb = torch.cat(lines, 2)*.5+.5 return to_pil(rgb) with gr.Blocks() as demo: gr.Markdown( """ # BerfScene: Bev-conditioned Equivariant Radiance Fields for Infinite 3D Scene Generation Qihang Zhang, Yinghao Xu, Yujun Shen, Bo Dai, Bolei Zhou*, Ceyuan Yang* (*Corresponding Author)
[Arxiv Report](https://arxiv.org/abs/2312.02136) | [Project Page](https://zqh0253.github.io/BerfScene/) | [Github](https://github.com/zqh0253/BerfScene) """ ) gr.Markdown( """ ### Quick Start 1. Drag and place objects in the canvas. 2. Click `Add object` to insert object into the canvas. 3. Click `Reset` to clean the canvas. 4. Click `Get local view` to synthesize local 3D scenes. 5. Click `Get global view` to synthesize global 3D scenes. """ ) with gr.Row(): with gr.Column(): drag = Draggable() with gr.Row(): submit_btn_local = gr.Button("Get local view", variant='primary') submit_btn_global = gr.Button("Get global view", variant='primary') with gr.Column(): with gr.Row(): single_view_image = gr.Image(label='single view', interactive=False) single_view_video = gr.Video(label='mutli-view', interactive=False, autoplay=True) global_view_image = gr.Image(label='global view', interactive=False) submit_btn_local.click(fn=predict_local_view, inputs=drag, outputs=single_view_image) submit_btn_local.click(fn=predict_local_view_video, inputs=drag, outputs=single_view_video) submit_btn_global.click(fn=predict_global_view, inputs=drag, outputs=global_view_image) parser = argparse.ArgumentParser() parser.add_argument('--port', type=int, help='The port number', default=7860) args = parser.parse_args() demo.queue() demo.launch(server_name='0.0.0.0', server_port=args.port, debug=True, show_error=True) # demo.launch(server_name='0.0.0.0', server_port=7860, debug=True, show_error=True)