yslan/ObjCtrl-2.5D
10
1try:2 import spaces3except:4 pass5 6import os7import gradio as gr8import json9import ast10 11import torch12from gradio_image_prompter import ImagePrompter13from sam2.sam2_image_predictor import SAM2ImagePredictor14from omegaconf import OmegaConf15from PIL import Image, ImageDraw16import numpy as np17from copy import deepcopy18import cv219 20import torch.nn.functional as F21import torchvision22from einops import rearrange23import tempfile24 25from objctrl_2_5d.utils.ui_utils import process_image, get_camera_pose, get_subject_points, get_points, undo_points, mask_image, traj2cam, get_mid_params26from ZoeDepth.zoedepth.utils.misc import colorize27 28from cameractrl.inference import get_pipeline29 30from objctrl_2_5d.utils.objmask_util import RT2Plucker, Unprojected, roll_with_ignore_multidim, dilate_mask_pytorch31from objctrl_2_5d.utils.filter_utils import get_freq_filter, freq_mix_3d32 33### Title and Description ###34#### Description ####35title = r"""<h1 align="center">ObjCtrl-2.5D: Training-free Object Control with Camera Poses</h1>"""36# subtitle = r"""<h2 align="center">Deployed on SVD Generation</h2>"""37important_link = r"""38<div align='center'>39 <a href='https://wzhouxiff.github.io/projects/ObjCtrl-2.5D/assets/ObjCtrl-2.5D.pdf'>[Paper]</a>40  <a href='https://arxiv.org/pdf/2412.07721'>[arxiv]</a>41  <a href='https://wzhouxiff.github.io/projects/ObjCtrl-2.5D/'>[Project Page]</a>42  <a href='https://github.com/wzhouxiff/ObjCtrl-2.5D'>[Code]</a>43</div>44"""45 46# authors = r"""47# <div align='center'>48# <a href='https://wzhouxiff.github.io/'>Zhouxia Wang</a>49#   <a href='https://nirvanalan.github.io/'>Yushi Lan</a>50#   <a href='https://shangchenzhou.com/'>Shangchen Zhou</a>51#   <a href='https://www.mmlab-ntu.com/person/ccloy/index.html'>Chen Change Loy</a>52# </div>53# """54 55# affiliation = r"""56# <div align='center'>57# <a href='https://www.mmlab-ntu.com/'>S-Lab, NTU Singapore</a>58# </div>59# """60 61description = r"""62<b>Official Gradio demo</b> for <a href='https://github.com/wzhouxiff/ObjCtrl-2.5D' target='_blank'><b>ObjCtrl-2.5D: Training-free Object Control with Camera Poses</b></a>.<br>63๐ฅ ObjCtrl2.5D enables object motion control in a I2V generated video via transforming 2D trajectories to 3D using depth, subsequently converting them into camera poses, 64thereby leveraging the exisitng camera motion control module for object motion control without requiring additional training.<br>65"""66 67article = r"""68If ObjCtrl2.5D is helpful, please help to โญ the <a href='https://github.com/wzhouxiff/ObjCtrl-2.5D' target='_blank'>Github Repo</a>. Thanks! 69[](https://github.com/wzhouxiff/ObjCtrl-2.5D)71 72---73 74๐ **License**75<br>76This project is licensed under <a href="https://github.com/wzhouxiff/ObjCtrl-2.5D/blob/main/LICENSE">S-Lab License 1.0</a>, 77Redistribution and use for non-commercial purposes should follow this license.78 79๐ **Citation**80<br>81If our work is useful for your research, please consider citing:82```bibtex83@inproceedings{objctrl2.5d,84 title={ObjCtrl-2.5D: Training-free Object Control with Camera Poses},85 author={Wang, Zhouxia and Lan, Yushi and Zhou, Shangchen and Loy, Chen Change},86 booktitle={arXiv preprint arXiv:2412.07721},87 year={2024}88}89```90 91๐ง **Contact**92<br>93If you have any questions, please feel free to reach me out at <b>zhouzi1212@gmail.com</b>.94 95"""96 97# pre-defined parameters98DEBUG = False99 100if DEBUG:101 cur_OUTPUT_PATH = 'outputs/tmp'102 os.makedirs(cur_OUTPUT_PATH, exist_ok=True)103 104# num_inference_steps=25105min_guidance_scale = 1.0106max_guidance_scale = 3.0107 108area_ratio = 0.3109depth_scale_ = 5.2110center_margin = 10111 112height, width = 320, 576113num_frames = 14114 115intrinsics = np.array([[float(width), float(width), float(width) / 2, float(height) / 2]])116intrinsics = np.repeat(intrinsics, num_frames, axis=0) # [n_frame, 4]117fx = intrinsics[0, 0] / width118fy = intrinsics[0, 1] / height119cx = intrinsics[0, 2] / width120cy = intrinsics[0, 3] / height121 122down_scale = 8123H, W = height // down_scale, width // down_scale124K = np.array([[width / down_scale, 0, W / 2], [0, width / down_scale, H / 2], [0, 0, 1]])125 126 127# -------------- initialization --------------128 129# CAMERA_MODE = ["Traj2Cam", "Rotate", "Clockwise", "Translate"]130CAMERA_MODE = ["None", "ZoomIn", "ZoomOut", "PanRight", "PanLeft", "TiltUp", "TiltDown", "ClockWise", "Anti-CW", "Rotate60"]131 132# select the device for computation133if torch.cuda.is_available():134 device = torch.device("cuda")135elif torch.backends.mps.is_available():136 device = torch.device("mps")137else:138 device = torch.device("cpu")139print(f"using device: {device}")140 141# # segmentation model142segmentor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-tiny", cache_dir="ckpt", device=device)143 144# depth model145d_model_NK = torch.hub.load('./ZoeDepth', 'ZoeD_NK', source='local', pretrained=True).to(device)146 147# cameractrl model148config = "configs/svd_320_576_cameractrl.yaml"149model_id = "stabilityai/stable-video-diffusion-img2vid"150ckpt = "checkpoints/CameraCtrl_svd.ckpt"151if not os.path.exists(ckpt):152 os.makedirs("checkpoints", exist_ok=True)153 os.system("wget -c https://huggingface.co/hehao13/CameraCtrl_SVD_ckpts/resolve/main/CameraCtrl_svd.ckpt?download=true")154 os.system("mv CameraCtrl_svd.ckpt?download=true checkpoints/CameraCtrl_svd.ckpt")155model_config = OmegaConf.load(config)156 157 158pipeline = get_pipeline(model_id, "unet", model_config['down_block_types'], model_config['up_block_types'],159 model_config['pose_encoder_kwargs'], model_config['attention_processor_kwargs'],160 ckpt, True, device)161 162# segmentor = None163# d_model_NK = None164# pipeline = None165 166### run the demo ##167@spaces.GPU(duration=7)168def segment(canvas, image, logits):169 if logits is not None:170 logits *= 32.0171 _, points = get_subject_points(canvas)172 image = np.array(image)173 174 with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):175 segmentor.set_image(image)176 input_points = []177 input_boxes = []178 for p in points:179 [x1, y1, _, x2, y2, _] = p180 if x2==0 and y2==0:181 input_points.append([x1, y1])182 else:183 input_boxes.append([x1, y1, x2, y2])184 if len(input_points) == 0:185 input_points = None186 input_labels = None187 else:188 input_points = np.array(input_points)189 input_labels = np.ones(len(input_points))190 if len(input_boxes) == 0:191 input_boxes = None192 else:193 input_boxes = np.array(input_boxes)194 masks, _, logits = segmentor.predict(195 point_coords=input_points,196 point_labels=input_labels,197 box=input_boxes,198 multimask_output=False,199 return_logits=True,200 mask_input=logits,201 )202 mask = masks > 0203 masked_img = mask_image(image, mask[0], color=[252, 140, 90], alpha=0.9)204 masked_img = Image.fromarray(masked_img)205 206 return mask[0], {'image': masked_img, 'points': points}, logits / 32.0207 208@spaces.GPU(duration=80)209def run_objctrl_2_5d(condition_image, 210 mask, 211 depth, 212 RTs, 213 bg_mode, 214 shared_wapring_latents, 215 scale_wise_masks, 216 rescale, 217 seed, 218 ds, dt, 219 num_inference_steps=25):220 seed = int(seed)221 222 center_h_margin, center_w_margin = center_margin, center_margin223 depth_center = np.mean(depth[height//2-center_h_margin:height//2+center_h_margin, width//2-center_w_margin:width//2+center_w_margin])224 225 if rescale > 0:226 depth_rescale = round(depth_scale_ * rescale / depth_center, 2)227 else:228 depth_rescale = 1.0229 230 depth = depth * depth_rescale231 232 depth_down = F.interpolate(torch.tensor(depth).unsqueeze(0).unsqueeze(0), 233 (H, W), mode='bilinear', align_corners=False).squeeze().numpy() # [H, W]234 235 ## latent236 generator = torch.Generator()237 generator.manual_seed(seed)238 239 latents_org = pipeline.prepare_latents(240 1,241 14,242 8,243 height,244 width,245 pipeline.dtype,246 device,247 generator,248 None,249 )250 latents_org = latents_org / pipeline.scheduler.init_noise_sigma251 252 cur_plucker_embedding, _, _ = RT2Plucker(RTs, RTs.shape[0], (height, width), fx, fy, cx, cy) # 6, V, H, W253 cur_plucker_embedding = cur_plucker_embedding.to(device)254 cur_plucker_embedding = cur_plucker_embedding[None, ...] # b 6 f h w255 cur_plucker_embedding = cur_plucker_embedding.permute(0, 2, 1, 3, 4) # b f 6 h w256 cur_plucker_embedding = cur_plucker_embedding[:, :num_frames, ...]257 cur_pose_features = pipeline.pose_encoder(cur_plucker_embedding)258 259 # bg_mode = ["Fixed", "Reverse", "Free"]260 if bg_mode == "Fixed":261 fix_RTs = np.repeat(RTs[0][None, ...], num_frames, axis=0) # [n_frame, 4, 3]262 fix_plucker_embedding, _, _ = RT2Plucker(fix_RTs, num_frames, (height, width), fx, fy, cx, cy) # 6, V, H, W263 fix_plucker_embedding = fix_plucker_embedding.to(device)264 fix_plucker_embedding = fix_plucker_embedding[None, ...] # b 6 f h w265 fix_plucker_embedding = fix_plucker_embedding.permute(0, 2, 1, 3, 4) # b f 6 h w266 fix_plucker_embedding = fix_plucker_embedding[:, :num_frames, ...]267 fix_pose_features = pipeline.pose_encoder(fix_plucker_embedding)268 269 elif bg_mode == "Reverse":270 bg_plucker_embedding, _, _ = RT2Plucker(RTs[::-1], RTs.shape[0], (height, width), fx, fy, cx, cy) # 6, V, H, W271 bg_plucker_embedding = bg_plucker_embedding.to(device)272 bg_plucker_embedding = bg_plucker_embedding[None, ...] # b 6 f h w273 bg_plucker_embedding = bg_plucker_embedding.permute(0, 2, 1, 3, 4) # b f 6 h w274 bg_plucker_embedding = bg_plucker_embedding[:, :num_frames, ...]275 fix_pose_features = pipeline.pose_encoder(bg_plucker_embedding)276 277 else:278 fix_pose_features = None279 280 #### preparing mask281 282 mask = Image.fromarray(mask)283 mask = mask.resize((W, H))284 mask = np.array(mask).astype(np.float32)285 mask = np.expand_dims(mask, axis=-1)286 287 # visulize mask288 if DEBUG:289 mask_sum_vis = mask[..., 0]290 mask_sum_vis = (mask_sum_vis * 255.0).astype(np.uint8)291 mask_sum_vis = Image.fromarray(mask_sum_vis)292 293 mask_sum_vis.save(f'{cur_OUTPUT_PATH}/org_mask.png')294 295 try:296 warped_masks = Unprojected(mask, depth_down, RTs, H=H, W=W, K=K)297 298 warped_masks.insert(0, mask)299 300 except:301 # mask to bbox302 print(f'!!! Mask is too small to warp; mask to bbox') 303 mask = mask[:, :, 0]304 coords = cv2.findNonZero(mask)305 x, y, w, h = cv2.boundingRect(coords)306 # mask[y:y+h, x:x+w] = 1.0307 308 center_x, center_y = x + w // 2, y + h // 2309 center_z = depth_down[center_y, center_x]310 311 # RTs [n_frame, 3, 4] to [n_frame, 4, 4] , add [0, 0, 0, 1]312 RTs = np.concatenate([RTs, np.array([[[0, 0, 0, 1]]] * num_frames)], axis=1)313 314 # RTs: world to camera315 P0 = np.array([center_x, center_y, 1])316 Pc0 = np.linalg.inv(K) @ P0 * center_z317 pw = np.linalg.inv(RTs[0]) @ np.array([Pc0[0], Pc0[1], center_z, 1]) # [4]318 319 P = [np.array([center_x, center_y])]320 for i in range(1, num_frames):321 Pci = RTs[i] @ pw322 Pi = K @ Pci[:3] / Pci[2]323 P.append(Pi[:2])324 325 warped_masks = [mask]326 for i in range(1, num_frames):327 shift_x = int(round(P[i][0] - P[0][0]))328 shift_y = int(round(P[i][1] - P[0][1]))329 330 cur_mask = roll_with_ignore_multidim(mask, [shift_y, shift_x])331 warped_masks.append(cur_mask)332 333 334 warped_masks = [v[..., None] for v in warped_masks]335 336 warped_masks = np.stack(warped_masks, axis=0) # [f, h, w]337 warped_masks = np.repeat(warped_masks, 3, axis=-1) # [f, h, w, 3]338 339 mask_sum = np.sum(warped_masks, axis=0, keepdims=True) # [1, H, W, 3]340 mask_sum[mask_sum > 1.0] = 1.0341 mask_sum = mask_sum[0,:,:, 0]342 343 if DEBUG:344 ## visulize warp mask 345 warp_masks_vis = torch.tensor(warped_masks)346 warp_masks_vis = (warp_masks_vis * 255.0).to(torch.uint8)347 torchvision.io.write_video(f'{cur_OUTPUT_PATH}/warped_masks.mp4', warp_masks_vis, fps=10, video_codec='h264', options={'crf': '10'})348 349 # visulize mask350 mask_sum_vis = mask_sum351 mask_sum_vis = (mask_sum_vis * 255.0).astype(np.uint8)352 mask_sum_vis = Image.fromarray(mask_sum_vis)353 354 mask_sum_vis.save(f'{cur_OUTPUT_PATH}/merged_mask.png')355 356 if scale_wise_masks:357 min_area = H * W * area_ratio # cal in downscale358 non_zero_len = mask_sum.sum() 359 360 print(f'non_zero_len: {non_zero_len}, min_area: {min_area}')361 362 if non_zero_len > min_area:363 kernel_sizes = [1, 1, 1, 3]364 elif non_zero_len > min_area * 0.5:365 kernel_sizes = [3, 1, 1, 5]366 else:367 kernel_sizes = [5, 3, 3, 7]368 else:369 kernel_sizes = [1, 1, 1, 1]370 371 mask = torch.from_numpy(mask_sum) # [h, w]372 mask = mask[None, None, ...] # [1, 1, h, w]373 mask = F.interpolate(mask, (height, width), mode='bilinear', align_corners=False) # [1, 1, H, W]374 # mask = mask.repeat(1, num_frames, 1, 1) # [1, f, H, W]375 mask = mask.to(pipeline.dtype).to(device)376 377 ##### Mask End ######378 379 ### Got blending pose features Start ###380 381 pose_features = []382 for i in range(0, len(cur_pose_features)):383 kernel_size = kernel_sizes[i]384 h, w = cur_pose_features[i].shape[-2:]385 386 if fix_pose_features is None:387 pose_features.append(torch.zeros_like(cur_pose_features[i]))388 else:389 pose_features.append(fix_pose_features[i])390 391 cur_mask = F.interpolate(mask, (h, w), mode='bilinear', align_corners=False)392 cur_mask = dilate_mask_pytorch(cur_mask, kernel_size=kernel_size) # [1, 1, H, W]393 cur_mask = cur_mask.repeat(1, num_frames, 1, 1) # [1, f, H, W]394 395 if DEBUG:396 # visulize mask397 mask_vis = cur_mask[0, 0].cpu().numpy() * 255.0398 mask_vis = Image.fromarray(mask_vis.astype(np.uint8))399 mask_vis.save(f'{cur_OUTPUT_PATH}/mask_k{kernel_size}_scale{i}.png')400 401 cur_mask = cur_mask[None, ...] # [1, 1, f, H, W]402 pose_features[-1] = cur_pose_features[i] * cur_mask + pose_features[-1] * (1 - cur_mask)403 404 ### Got blending pose features End ###405 406 ##### Warp Noise Start ######407 408 if shared_wapring_latents:409 noise = latents_org[0, 0].data.cpu().numpy().copy() #[14, 4, 40, 72]410 noise = np.transpose(noise, (1, 2, 0)) # [40, 72, 4]411 412 try:413 warp_noise = Unprojected(noise, depth_down, RTs, H=H, W=W, K=K)414 warp_noise.insert(0, noise)415 except:416 print(f'!!! Noise is too small to warp; mask to bbox')417 418 warp_noise = [noise]419 for i in range(1, num_frames):420 shift_x = int(round(P[i][0] - P[0][0]))421 shift_y = int(round(P[i][1] - P[0][1]))422 423 cur_noise= roll_with_ignore_multidim(noise, [shift_y, shift_x])424 warp_noise.append(cur_noise)425 426 warp_noise = np.stack(warp_noise, axis=0) # [f, h, w, 4]427 428 if DEBUG:429 ## visulize warp noise430 warp_noise_vis = torch.tensor(warp_noise)[..., :3] * torch.tensor(warped_masks)431 warp_noise_vis = (warp_noise_vis - warp_noise_vis.min()) / (warp_noise_vis.max() - warp_noise_vis.min())432 warp_noise_vis = (warp_noise_vis * 255.0).to(torch.uint8)433 434 torchvision.io.write_video(f'{cur_OUTPUT_PATH}/warp_noise.mp4', warp_noise_vis, fps=10, video_codec='h264', options={'crf': '10'})435 436 437 warp_latents = torch.tensor(warp_noise).permute(0, 3, 1, 2).to(latents_org.device).to(latents_org.dtype) # [frame, 4, H, W]438 warp_latents = warp_latents.unsqueeze(0) # [1, frame, 4, H, W]439 440 warped_masks = torch.tensor(warped_masks).permute(0, 3, 1, 2).unsqueeze(0) # [1, frame, 3, H, W]441 mask_extend = torch.concat([warped_masks, warped_masks[:,:,0:1]], dim=2) # [1, frame, 4, H, W]442 mask_extend = mask_extend.to(latents_org.device).to(latents_org.dtype)443 444 warp_latents = warp_latents * mask_extend + latents_org * (1 - mask_extend)445 warp_latents = warp_latents.permute(0, 2, 1, 3, 4)446 random_noise = latents_org.clone().permute(0, 2, 1, 3, 4)447 448 filter_shape = warp_latents.shape449 450 freq_filter = get_freq_filter(451 filter_shape, 452 device = device, 453 filter_type='butterworth',454 n=4,455 d_s=ds,456 d_t=dt457 )458 459 warp_latents = freq_mix_3d(warp_latents, random_noise, freq_filter)460 warp_latents = warp_latents.permute(0, 2, 1, 3, 4)461 462 else:463 warp_latents = latents_org.clone()464 465 generator.manual_seed(42)466 467 with torch.no_grad():468 result = pipeline(469 image=condition_image,470 pose_embedding=cur_plucker_embedding,471 height=height,472 width=width,473 num_frames=num_frames,474 num_inference_steps=num_inference_steps,475 min_guidance_scale=min_guidance_scale,476 max_guidance_scale=max_guidance_scale,477 do_image_process=True,478 generator=generator,479 output_type='pt',480 pose_features= pose_features,481 latents = warp_latents482 ).frames[0].cpu() #[f, c, h, w]483 484 485 result = rearrange(result, 'f c h w -> f h w c')486 result = (result * 255.0).to(torch.uint8)487 488 video_path = tempfile.NamedTemporaryFile(suffix='.mp4').name489 torchvision.io.write_video(video_path, result, fps=10, video_codec='h264', options={'crf': '8'})490 491 return video_path492 493 494# UI function495@spaces.GPU(duration=7)496def process_image(raw_image, trajectory_points):497 498 image, points = raw_image['image'], raw_image['points']499 500 print(points)501 502 try:503 assert(len(points)) == 1, "Please draw only one bbox"504 [x1, y1, _, x2, y2, _] = points[0]505 506 image = image.crop((x1, y1, x2, y2))507 image = image.resize((width, height))508 except:509 image = image.resize((width, height))510 511 depth = d_model_NK.infer_pil(image) 512 colored_depth = colorize(depth, cmap='gray_r') # [h, w, 4] 0-255513 514 depth_img = deepcopy(colored_depth[:, :, :3])515 if len(trajectory_points) > 0:516 for idx, point in enumerate(trajectory_points):517 if idx % 2 == 0:518 cv2.circle(depth_img, tuple(point), 10, (255, 0, 0), -1)519 else:520 cv2.circle(depth_img, tuple(point), 10, (0, 0, 255), -1)521 if idx > 0:522 line_length = np.sqrt((trajectory_points[idx][0] - trajectory_points[idx-1][0])**2 + (trajectory_points[idx][1] - trajectory_points[idx-1][1])**2)523 arrow_head_length = 10524 tip_length = arrow_head_length / line_length525 cv2.arrowedLine(depth_img, trajectory_points[idx-1], trajectory_points[idx], (0, 255, 0), 4, tipLength=tip_length)526 527 return image, {'image': image}, depth, depth_img, colored_depth[:, :, :3]528 529 530 531def draw_points_on_image(img, points):532 # img = Image.fromarray(np.array(image))533 draw = ImageDraw.Draw(img)534 535 for p in points:536 x1, y1, _, x2, y2, _ = p537 538 if x2 == 0 and y2 == 0:539 # Point: ้่ฒ็นๅธฆ้ป่พน540 point_radius = 4541 draw.ellipse(542 (x1 - point_radius, y1 - point_radius, x1 + point_radius, y1 + point_radius),543 fill="cyan", outline="black", width=1544 )545 else:546 # Bounding Box: ้ป่ฒ็ฉๅฝขๆก547 draw.rectangle([x1, y1, x2, y2], outline="black", width=3)548 549 return img550 551@spaces.GPU(duration=15)552def from_examples(raw_input, raw_image_points, canvas, seg_image_points, selected_points_text, camera_option, mask_bk):553 raw_image_points = ast.literal_eval(raw_image_points)554 seg_image_points = ast.literal_eval(seg_image_points)555 556 selected_points = ast.literal_eval(selected_points_text)557 mask = np.array(mask_bk)558 mask = mask[:,:,0] > 0559 selected_points = ast.literal_eval(selected_points_text)560 561 image, _, depth, depth_img, colored_depth = process_image({'image': raw_input['image'], 'points': raw_image_points}, selected_points)562 563 # get camera pose564 if camera_option == "None":565 # traj2came566 rescale = 1.0567 camera_pose, camera_pose_vis, rescale, _ = traj2cam(selected_points, depth , rescale)568 else:569 rescale = 0.0570 angle = 60571 speed = 4.0572 camera_pose, camera_pose_vis, rescale = get_camera_pose(CAMERA_MODE)(camera_option, depth, mask, rescale, angle, speed)573 574 575 raw_image = draw_points_on_image(raw_input['image'], raw_image_points)576 seg_image = draw_points_on_image(canvas['image'], seg_image_points)577 578 return image, mask, depth, depth_img, colored_depth, camera_pose, \579 camera_pose_vis, rescale, selected_points, \580 gr.update(value={'image': raw_image, 'points': raw_image_points}), \581 gr.update(value={'image': seg_image, 'points': seg_image_points}), \582 583 584# -------------- UI definition --------------585with gr.Blocks() as demo:586 # layout definition587 gr.Markdown(title)588 # gr.Markdown(authors)589 # gr.Markdown(affiliation)590 gr.Markdown(important_link)591 gr.Markdown(description)592 593 594 # with gr.Row():595 # gr.Markdown("""# <center>Repositioning the Subject within Image </center>""")596 mask = gr.State(value=None) # store mask597 mask_bk = gr.Image(type="pil", label="Mask", show_label=True, interactive=False, visible=False)598 599 removal_mask = gr.State(value=None) # store removal mask600 selected_points = gr.State([]) # store points601 selected_points_text = gr.Textbox(label="Selected Points", visible=False)602 raw_image_points = gr.Textbox(label="Raw Image Points", visible=False)603 seg_image_points = gr.Textbox(label="Segment Image Points", visible=False)604 605 original_image = gr.State(value=None) # store original input image606 # masked_original_image = gr.State(value=None) # store masked input image607 mask_logits = gr.State(value=None) # store mask logits608 609 depth = gr.State(value=None) # store depth610 org_depth_image = gr.State(value=None) # store original depth image611 612 camera_pose = gr.State(value=None) # store camera pose613 614 rescale = gr.Slider(minimum=0.0, maximum=10, step=0.1, value=1.0, label="Rescale", interactive=True, visible=False)615 angle = gr.Slider(minimum=-360, maximum=360, step=1, value=60, label="Angle", interactive=True, visible=False)616 617 seed = gr.Textbox(value = "42", label="Seed", interactive=True, visible=False)618 scale_wise_masks = gr.Checkbox(label="Enable Scale-wise Masks", interactive=True, value=True, visible=False)619 ds = gr.Slider(minimum=0.0, maximum=1, step=0.1, value=0.25, label="ds", interactive=True, visible=False)620 dt = gr.Slider(minimum=0.0, maximum=1, step=0.1, value=0.1, label="dt", interactive=True, visible=False)621 622 with gr.Column():623 624 outlines = """625 <font size="5"><b>There are total 5 steps to complete the task.</b></font>626 - Step 1: Input an image and Crop it to a suitable size and attained depth;627 - Step 2: Attain the subject mask;628 - Step 3: Draw trajectory on depth map or skip to use camera pose;629 - Step 4: Select camera poses or skip.630 - Step 5: Generate the final video.631 """632 633 gr.Markdown(outlines)634 635 636 with gr.Row():637 with gr.Column():638 # Step 1: Input Image639 step1_dec = """640 <font size="4"><b>Step 1: Input Image</b></font>641 """642 step1 = gr.Markdown(step1_dec)643 raw_input = ImagePrompter(type="pil", label="Raw Image", show_label=True, interactive=True)644 645 step1_notes = """646 - Select the region using a <mark>bounding box</mark>, aiming for a ratio close to </mark>320:576</mark> (height:width).647 - If the input is in 320 x 576, press `Process` directly.648 """649 notes = gr.Markdown(step1_notes)650 651 process_button = gr.Button("Process")652 653 with gr.Column():654 # Step 2: Get Subject Mask655 step2_dec = """656 <font size="4"><b>Step 2: Get Subject Mask</b></font>657 """658 step2 = gr.Markdown(step2_dec)659 canvas = ImagePrompter(type="pil", label="Input Image", show_label=True, interactive=True) # for mask painting660 661 step2_notes = """662 - Use the <mark>bounding boxes</mark> or <mark>points</mark> to select the subject.663 - Press `Segment Subject` to get the mask. <mark>Can be refined iteratively by updating points<mark>.664 """665 notes = gr.Markdown(step2_notes)666 667 select_button = gr.Button("Segment Subject")668 669 with gr.Column():670 # Step 3: Get Depth and Draw Trajectory671 step3_dec = """672 <font size="4"><b>Step 3: Draw Trajectory on Depth or <mark>SKIP</mark></b></font>673 674 """675 step3 = gr.Markdown(step3_dec)676 depth_image = gr.Image(type="pil", label="Depth Image", show_label=True, interactive=False)677 678 step3_dec = """679 - Selecting points on the depth image. <mark>No more than 14 points</mark>.680 - Press `Undo point` to remove all points. Press `Traj2Cam` to get camera poses.681 """682 notes = gr.Markdown(step3_dec)683 684 undo_button = gr.Button("Undo point")685 traj2cam_button = gr.Button("Traj2Cam")686 687 with gr.Row():688 689 with gr.Column():690 # Step 4: Trajectory to Camera Pose or Get Camera Pose691 step4_dec = """692 <font size="4"><b>Step 4: Get Customized Camera Poses or <mark>SKIP</mark></b></font>693 """694 step4 = gr.Markdown(step4_dec)695 camera_pose_vis = gr.Plot(None, label='Camera Pose')696 camera_option = gr.Radio(choices = CAMERA_MODE, label='Camera Options', value=CAMERA_MODE[0], interactive=True)697 speed = gr.Slider(minimum=0.1, maximum=10, step=0.1, value=4.0, label="Speed", interactive=True, visible=True)698 699 with gr.Column():700 # Step 5: Get the final generated video701 step5_dec = """702 <font size="4"><b>Step 5: Get the Final Generated Video</b></font>703 """704 step5 = gr.Markdown(step5_dec)705 generated_video = gr.Video(None, label='Generated Video')706 707 # with gr.Row():708 bg_mode = gr.Radio(choices = ["Fixed", "Reverse", "Free"], label="Background Mode", value="Fixed", interactive=True)709 shared_wapring_latents = gr.Checkbox(label="Enable Shared Warping Latents", interactive=True, value=False, visible=True)710 711 generated_button = gr.Button("Generate")712 713 get_mid_params_button = gr.Button("Get Mid Params", visible=False)714 715 716 # # event definition717 process_button.click(718 fn = process_image,719 inputs = [raw_input, selected_points],720 outputs = [original_image, canvas, depth, depth_image, org_depth_image]721 )722 723 select_button.click(724 segment,725 [canvas, original_image, mask_logits],726 [mask, canvas, mask_logits]727 )728 729 depth_image.select(730 get_points,731 [depth_image, selected_points],732 [depth_image, selected_points],733 )734 undo_button.click(735 undo_points,736 [org_depth_image],737 [depth_image, selected_points]738 )739 740 traj2cam_button.click(741 traj2cam,742 [selected_points, depth, rescale],743 [camera_pose, camera_pose_vis, rescale, camera_option]744 )745 746 camera_option.change(747 get_camera_pose(CAMERA_MODE),748 [camera_option, depth, mask, rescale, angle, speed],749 [camera_pose, camera_pose_vis, rescale]750 )751 752 generated_button.click(753 run_objctrl_2_5d,754 [755 original_image,756 mask,757 depth,758 camera_pose,759 bg_mode,760 shared_wapring_latents,761 scale_wise_masks,762 rescale,763 seed,764 ds,765 dt,766 # num_inference_steps767 ],768 [generated_video],769 )770 771 get_mid_params_button.click(772 get_mid_params,773 [raw_input, canvas, mask, selected_points, camera_option, bg_mode, shared_wapring_latents, generated_video]774 )775 776 ## Get examples777 with open('./assets/examples/examples.json', 'r') as f:778 examples = json.load(f)779 # print(examples)780 781 # examples = [examples]782 examples = [v for k, v in examples.items()]783 784 gr.Examples(785 examples=examples,786 inputs=[787 raw_input,788 raw_image_points,789 canvas,790 seg_image_points,791 mask_bk,792 selected_points_text, # selected_points793 camera_option,794 bg_mode,795 shared_wapring_latents,796 generated_video797 ],798 examples_per_page=20799 )800 801 selected_points_text.change(802 from_examples,803 inputs=[raw_input, raw_image_points, canvas, seg_image_points, selected_points_text, camera_option, mask_bk],804 outputs=[original_image, mask, depth, depth_image, org_depth_image, camera_pose, camera_pose_vis, rescale, selected_points, raw_input, canvas]805 )806 807 808 809 810 gr.Markdown(article)811 812 813demo.queue().launch(share=True)814