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 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 # 使用OpenCV创建视频写入对象，选择H.264编码器 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()