import gradio as gr
import numpy as np
import os
from PIL import Image
import torch
import torchvision.transforms as transforms
import options
import test
import importlib
from scipy.interpolate import interp1d, splev, splprep
import cv2
import subprocess
subprocess.run(["bash", "install_imaginaire.sh"])
def get_single(sat_img, style_img, x_offset, y_offset):
name = ''
for i in [name for name in os.listdir('demo_img') if 'case' in name]:
style = Image.open('demo_img/{}/groundview.image.png'.format(i)).convert('RGB')
style =np.array(style)
if (style == style_img).all():
name = i
break
input_dict = {}
trans = transforms.ToTensor()
input_dict['sat'] = trans(sat_img)
input_dict['pano'] = trans(style_img)
input_dict['paths'] = "demo.png"
sky = trans(Image.open('demo_img/{}/groundview.sky.png'.format(name)).convert("L"))
input_a = input_dict['pano']*sky
sky_histc = torch.cat([input_a[i].histc()[10:] for i in reversed(range(3))])
input_dict['sky_histc'] = sky_histc
input_dict['sky_mask'] = sky
for key in input_dict.keys():
if isinstance(input_dict[key], torch.Tensor):
input_dict[key] = input_dict[key].unsqueeze(0)
args = ["--yaml=sat2density_cvact", "--test_ckpt_path=wandb/run-20230219_141512-2u87bj8w/files/checkpoint/model.pth", "--task=test_vid", "--demo_img=demo_img/case1/satview-input.png",
"--sty_img=demo_img/case1/groundview.image.png", "--save_dir=output"]
opt_cmd = options.parse_arguments(args=args)
opt = options.set(opt_cmd=opt_cmd)
opt.isTrain = False
opt.name = opt.yaml if opt.name is None else opt.name
opt.batch_size = 1
m = importlib.import_module("model.{}".format(opt.model))
model = m.Model(opt)
# m.load_dataset(opt)
model.build_networks(opt)
ckpt = torch.load(opt.test_ckpt_path, map_location='cpu')
model.netG.load_state_dict(ckpt['netG'])
model.netG.eval()
model.set_input(input_dict)
model.style_temp = model.sky_histc
opt.origin_H_W = [-(y_offset*256-128)/128, (x_offset*256-128)/128] # TODO: hard code should be removed in the future
model.forward(opt)
rgb = model.out_put.pred[0].clamp(min=0,max=1.0).cpu().detach().numpy().transpose((1,2,0))
rgb = np.array(rgb*255, dtype=np.uint8)
return rgb
def get_video(sat_img, style_img, positions):
name = ''
for i in [name for name in os.listdir('demo_img') if 'case' in name]:
style = Image.open('demo_img/{}/groundview.image.png'.format(i)).convert('RGB')
style =np.array(style)
if (style == style_img).all():
name = i
break
input_dict = {}
trans = transforms.ToTensor()
input_dict['sat'] = trans(sat_img)
input_dict['pano'] = trans(style_img)
input_dict['paths'] = "demo.png"
sky = trans(Image.open('demo_img/{}/groundview.sky.png'.format(name)).convert("L"))
input_a = input_dict['pano']*sky
sky_histc = torch.cat([input_a[i].histc()[10:] for i in reversed(range(3))])
input_dict['sky_histc'] = sky_histc
input_dict['sky_mask'] = sky
for key in input_dict.keys():
if isinstance(input_dict[key], torch.Tensor):
input_dict[key] = input_dict[key].unsqueeze(0)
args = ["--yaml=sat2density_cvact", "--test_ckpt_path=wandb/run-20230219_141512-2u87bj8w/files/checkpoint/model.pth", "--task=test_vid", "--demo_img=demo_img/case1/satview-input.png",
"--sty_img=demo_img/case1/groundview.image.png", "--save_dir=output"]
opt_cmd = options.parse_arguments(args=args)
opt = options.set(opt_cmd=opt_cmd)
opt.isTrain = False
opt.name = opt.yaml if opt.name is None else opt.name
opt.batch_size = 1
m = importlib.import_module("model.{}".format(opt.model))
model = m.Model(opt)
# m.load_dataset(opt)
model.build_networks(opt)
ckpt = torch.load(opt.test_ckpt_path, map_location='cpu')
model.netG.load_state_dict(ckpt['netG'])
model.netG.eval()
model.set_input(input_dict)
model.style_temp = model.sky_histc
unique_lst = list(dict.fromkeys(positions))
pixels = []
for x in positions:
if x in unique_lst:
if x not in pixels:
pixels.append(x)
pixels = np.array(pixels)
tck, u = splprep(pixels.T, s=25, per=0)
u_new = np.linspace(u.min(), u.max(), 80)
x_new, y_new = splev(u_new, tck)
smooth_path = np.array([x_new,y_new]).T
rendered_image_list = []
rendered_depth_list = []
for i, (x,y) in enumerate(smooth_path):
opt.origin_H_W = [(y-128)/128, (x-128)/128] # TODO: hard code should be removed in the future
print('Rendering at ({}, {})'.format(x,y))
model.forward(opt)
rgb = model.out_put.pred[0].clamp(min=0,max=1.0).cpu().detach().numpy().transpose((1,2,0))
rgb = np.array(rgb*255, dtype=np.uint8)
rendered_image_list.append(rgb)
rendered_depth_list.append(
model.out_put.depth[0,0].cpu().detach().numpy()
)
output_video_path = 'output_video.mp4'
frame_rate = 15
frame_width = 512
frame_height = 128
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_video_path, fourcc, frame_rate, (frame_width, frame_height))
for image_np in rendered_image_list:
image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
out.write(image_np)
out.release()
return "output_video.mp4"
def copy_image(image):
return image
def show_image_and_point(image, x, y):
x = int(x*image.shape[1])
y = image.shape[0]-int(y*image.shape[0])
mask = np.zeros(image.shape[:2])
radius = min(image.shape[0], image.shape[1])//60
for i in range(x-radius-2, x+radius+2):
for j in range(y-radius-2, y+radius+2):
if (i-x)**2+(j-y)**2<=radius**2:
mask[j, i] = 1
return (image, [(mask, 'render point')])
def add_select_point(image, evt: gr.SelectData, state1):
if state1 == None:
state1 = []
x, y = evt.index
state1.append((x, y))
print(state1)
radius = min(image.shape[0], image.shape[1])//60
for i in range(x-radius-2, x+radius+2):
for j in range(y-radius-2, y+radius+2):
if (i-x)**2+(j-y)**2<=radius**2:
image[j, i, :] = 0
return image, state1
def reset_select_points(image):
return image, []
with gr.Blocks() as demo:
gr.Markdown("# Sat2Density Demos")
gr.Markdown("### select/upload the satllite image and select the style image")
with gr.Row():
with gr.Column():
sat_img = gr.Image(source='upload', shape=[256, 256], interactive=True)
img_examples = gr.Examples(examples=['demo_img/{}/satview-input.png'.format(i) for i in os.listdir('demo_img') if 'case' in i],
inputs=sat_img, outputs=None, examples_per_page=20)
with gr.Column():
style_img = gr.Image()
style_examples = gr.Examples(examples=['demo_img/{}/groundview.image.png'.format(i) for i in os.listdir('demo_img') if 'case' in i],
inputs=style_img, outputs=None, examples_per_page=20)
gr.Markdown("### select a certain point to generate single groundview image")
with gr.Row():
with gr.Column():
with gr.Row():
with gr.Column():
slider_x = gr.Slider(0.2, 0.8, 0.5, label="x-axis position")
slider_y = gr.Slider(0.2, 0.8, 0.5, label="y-axis position")
btn_single = gr.Button(label="demo1")
annotation_image = gr.AnnotatedImage()
out_single = gr.Image()
gr.Markdown("### draw a trajectory on the map to generate video")
state_select_points = gr.State()
with gr.Row():
with gr.Column():
draw_img = gr.Image(shape=[256, 256], interactive=True)
with gr.Column():
out_video = gr.Video()
reset_btn =gr.Button(value="Reset")
btn_video = gr.Button(label="demo1")
sat_img.change(copy_image, inputs = sat_img, outputs=draw_img)
draw_img.select(add_select_point, [draw_img, state_select_points], [draw_img, state_select_points])
sat_img.change(show_image_and_point, inputs = [sat_img, slider_x, slider_y], outputs = annotation_image)
slider_x.change(show_image_and_point, inputs = [sat_img, slider_x, slider_y], outputs = annotation_image, show_progress='hidden')
slider_y.change(show_image_and_point, inputs = [sat_img, slider_x, slider_y], outputs = annotation_image, show_progress='hidden')
btn_single.click(get_single, inputs = [sat_img, style_img, slider_x, slider_y], outputs=out_single)
reset_btn.click(reset_select_points, [sat_img], [draw_img, state_select_points])
btn_video.click(get_video, inputs=[sat_img, style_img, state_select_points], outputs=out_video) # 触发
demo.launch()