yslan/ObjCtrl-2.5D
10
1import spaces2 3import os4import torch5import numpy as np6import torch.nn.functional as F7import cv28import torchvision9from PIL import Image10from einops import rearrange11import tempfile12 13from objctrl_2_5d.utils.objmask_util import RT2Plucker, Unprojected, roll_with_ignore_multidim, dilate_mask_pytorch14from objctrl_2_5d.utils.filter_utils import get_freq_filter, freq_mix_3d15 16DEBUG = False17 18if DEBUG:19 cur_OUTPUT_PATH = 'outputs/tmp'20 os.makedirs(cur_OUTPUT_PATH, exist_ok=True)21 22# num_inference_steps=2523min_guidance_scale = 1.024max_guidance_scale = 3.025 26area_ratio = 0.327depth_scale_ = 5.228center_margin = 1029 30height, width = 320, 57631num_frames = 1432 33intrinsics = np.array([[float(width), float(width), float(width) / 2, float(height) / 2]])34intrinsics = np.repeat(intrinsics, num_frames, axis=0) # [n_frame, 4]35fx = intrinsics[0, 0] / width36fy = intrinsics[0, 1] / height37cx = intrinsics[0, 2] / width38cy = intrinsics[0, 3] / height39 40down_scale = 841H, W = height // down_scale, width // down_scale42K = np.array([[width / down_scale, 0, W / 2], [0, width / down_scale, H / 2], [0, 0, 1]])43 44@spaces.GPU(duration=50)45def run(pipeline, device):46 def run_objctrl_2_5d(condition_image, 47 mask, 48 depth, 49 RTs, 50 bg_mode, 51 shared_wapring_latents, 52 scale_wise_masks, 53 rescale, 54 seed, 55 ds, dt, 56 num_inference_steps=25):57 58 seed = int(seed)59 60 center_h_margin, center_w_margin = center_margin, center_margin61 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])62 63 if rescale > 0:64 depth_rescale = round(depth_scale_ * rescale / depth_center, 2)65 else:66 depth_rescale = 1.067 68 depth = depth * depth_rescale69 70 depth_down = F.interpolate(torch.tensor(depth).unsqueeze(0).unsqueeze(0), 71 (H, W), mode='bilinear', align_corners=False).squeeze().numpy() # [H, W]72 73 ## latent74 generator = torch.Generator()75 generator.manual_seed(seed)76 77 latents_org = pipeline.prepare_latents(78 1,79 14,80 8,81 height,82 width,83 pipeline.dtype,84 device,85 generator,86 None,87 )88 latents_org = latents_org / pipeline.scheduler.init_noise_sigma89 90 cur_plucker_embedding, _, _ = RT2Plucker(RTs, RTs.shape[0], (height, width), fx, fy, cx, cy) # 6, V, H, W91 cur_plucker_embedding = cur_plucker_embedding.to(device)92 cur_plucker_embedding = cur_plucker_embedding[None, ...] # b 6 f h w93 cur_plucker_embedding = cur_plucker_embedding.permute(0, 2, 1, 3, 4) # b f 6 h w94 cur_plucker_embedding = cur_plucker_embedding[:, :num_frames, ...]95 cur_pose_features = pipeline.pose_encoder(cur_plucker_embedding)96 97 # bg_mode = ["Fixed", "Reverse", "Free"]98 if bg_mode == "Fixed":99 fix_RTs = np.repeat(RTs[0][None, ...], num_frames, axis=0) # [n_frame, 4, 3]100 fix_plucker_embedding, _, _ = RT2Plucker(fix_RTs, num_frames, (height, width), fx, fy, cx, cy) # 6, V, H, W101 fix_plucker_embedding = fix_plucker_embedding.to(device)102 fix_plucker_embedding = fix_plucker_embedding[None, ...] # b 6 f h w103 fix_plucker_embedding = fix_plucker_embedding.permute(0, 2, 1, 3, 4) # b f 6 h w104 fix_plucker_embedding = fix_plucker_embedding[:, :num_frames, ...]105 fix_pose_features = pipeline.pose_encoder(fix_plucker_embedding)106 107 elif bg_mode == "Reverse":108 bg_plucker_embedding, _, _ = RT2Plucker(RTs[::-1], RTs.shape[0], (height, width), fx, fy, cx, cy) # 6, V, H, W109 bg_plucker_embedding = bg_plucker_embedding.to(device)110 bg_plucker_embedding = bg_plucker_embedding[None, ...] # b 6 f h w111 bg_plucker_embedding = bg_plucker_embedding.permute(0, 2, 1, 3, 4) # b f 6 h w112 bg_plucker_embedding = bg_plucker_embedding[:, :num_frames, ...]113 fix_pose_features = pipeline.pose_encoder(bg_plucker_embedding)114 115 else:116 fix_pose_features = None117 118 #### preparing mask119 120 mask = Image.fromarray(mask)121 mask = mask.resize((W, H))122 mask = np.array(mask).astype(np.float32)123 mask = np.expand_dims(mask, axis=-1)124 125 # visulize mask126 if DEBUG:127 mask_sum_vis = mask[..., 0]128 mask_sum_vis = (mask_sum_vis * 255.0).astype(np.uint8)129 mask_sum_vis = Image.fromarray(mask_sum_vis)130 131 mask_sum_vis.save(f'{cur_OUTPUT_PATH}/org_mask.png')132 133 try:134 warped_masks = Unprojected(mask, depth_down, RTs, H=H, W=W, K=K)135 136 warped_masks.insert(0, mask)137 138 except:139 # mask to bbox140 print(f'!!! Mask is too small to warp; mask to bbox') 141 mask = mask[:, :, 0]142 coords = cv2.findNonZero(mask)143 x, y, w, h = cv2.boundingRect(coords)144 # mask[y:y+h, x:x+w] = 1.0145 146 center_x, center_y = x + w // 2, y + h // 2147 center_z = depth_down[center_y, center_x]148 149 # RTs [n_frame, 3, 4] to [n_frame, 4, 4] , add [0, 0, 0, 1]150 RTs = np.concatenate([RTs, np.array([[[0, 0, 0, 1]]] * num_frames)], axis=1)151 152 # RTs: world to camera153 P0 = np.array([center_x, center_y, 1])154 Pc0 = np.linalg.inv(K) @ P0 * center_z155 pw = np.linalg.inv(RTs[0]) @ np.array([Pc0[0], Pc0[1], center_z, 1]) # [4]156 157 P = [np.array([center_x, center_y])]158 for i in range(1, num_frames):159 Pci = RTs[i] @ pw160 Pi = K @ Pci[:3] / Pci[2]161 P.append(Pi[:2])162 163 warped_masks = [mask]164 for i in range(1, num_frames):165 shift_x = int(round(P[i][0] - P[0][0]))166 shift_y = int(round(P[i][1] - P[0][1]))167 168 cur_mask = roll_with_ignore_multidim(mask, [shift_y, shift_x])169 warped_masks.append(cur_mask)170 171 172 warped_masks = [v[..., None] for v in warped_masks]173 174 warped_masks = np.stack(warped_masks, axis=0) # [f, h, w]175 warped_masks = np.repeat(warped_masks, 3, axis=-1) # [f, h, w, 3]176 177 mask_sum = np.sum(warped_masks, axis=0, keepdims=True) # [1, H, W, 3]178 mask_sum[mask_sum > 1.0] = 1.0179 mask_sum = mask_sum[0,:,:, 0]180 181 if DEBUG:182 ## visulize warp mask 183 warp_masks_vis = torch.tensor(warped_masks)184 warp_masks_vis = (warp_masks_vis * 255.0).to(torch.uint8)185 torchvision.io.write_video(f'{cur_OUTPUT_PATH}/warped_masks.mp4', warp_masks_vis, fps=10, video_codec='h264', options={'crf': '10'})186 187 # visulize mask188 mask_sum_vis = mask_sum189 mask_sum_vis = (mask_sum_vis * 255.0).astype(np.uint8)190 mask_sum_vis = Image.fromarray(mask_sum_vis)191 192 mask_sum_vis.save(f'{cur_OUTPUT_PATH}/merged_mask.png')193 194 if scale_wise_masks:195 min_area = H * W * area_ratio # cal in downscale196 non_zero_len = mask_sum.sum() 197 198 print(f'non_zero_len: {non_zero_len}, min_area: {min_area}')199 200 if non_zero_len > min_area:201 kernel_sizes = [1, 1, 1, 3]202 elif non_zero_len > min_area * 0.5:203 kernel_sizes = [3, 1, 1, 5]204 else:205 kernel_sizes = [5, 3, 3, 7]206 else:207 kernel_sizes = [1, 1, 1, 1]208 209 mask = torch.from_numpy(mask_sum) # [h, w]210 mask = mask[None, None, ...] # [1, 1, h, w]211 mask = F.interpolate(mask, (height, width), mode='bilinear', align_corners=False) # [1, 1, H, W]212 # mask = mask.repeat(1, num_frames, 1, 1) # [1, f, H, W]213 mask = mask.to(pipeline.dtype).to(device)214 215 ##### Mask End ######216 217 ### Got blending pose features Start ###218 219 pose_features = []220 for i in range(0, len(cur_pose_features)):221 kernel_size = kernel_sizes[i]222 h, w = cur_pose_features[i].shape[-2:]223 224 if fix_pose_features is None:225 pose_features.append(torch.zeros_like(cur_pose_features[i]))226 else:227 pose_features.append(fix_pose_features[i])228 229 cur_mask = F.interpolate(mask, (h, w), mode='bilinear', align_corners=False)230 cur_mask = dilate_mask_pytorch(cur_mask, kernel_size=kernel_size) # [1, 1, H, W]231 cur_mask = cur_mask.repeat(1, num_frames, 1, 1) # [1, f, H, W]232 233 if DEBUG:234 # visulize mask235 mask_vis = cur_mask[0, 0].cpu().numpy() * 255.0236 mask_vis = Image.fromarray(mask_vis.astype(np.uint8))237 mask_vis.save(f'{cur_OUTPUT_PATH}/mask_k{kernel_size}_scale{i}.png')238 239 cur_mask = cur_mask[None, ...] # [1, 1, f, H, W]240 pose_features[-1] = cur_pose_features[i] * cur_mask + pose_features[-1] * (1 - cur_mask)241 242 ### Got blending pose features End ###243 244 ##### Warp Noise Start ######245 246 if shared_wapring_latents:247 noise = latents_org[0, 0].data.cpu().numpy().copy() #[14, 4, 40, 72]248 noise = np.transpose(noise, (1, 2, 0)) # [40, 72, 4]249 250 try:251 warp_noise = Unprojected(noise, depth_down, RTs, H=H, W=W, K=K)252 warp_noise.insert(0, noise)253 except:254 print(f'!!! Noise is too small to warp; mask to bbox')255 256 warp_noise = [noise]257 for i in range(1, num_frames):258 shift_x = int(round(P[i][0] - P[0][0]))259 shift_y = int(round(P[i][1] - P[0][1]))260 261 cur_noise= roll_with_ignore_multidim(noise, [shift_y, shift_x])262 warp_noise.append(cur_noise)263 264 warp_noise = np.stack(warp_noise, axis=0) # [f, h, w, 4]265 266 if DEBUG:267 ## visulize warp noise268 warp_noise_vis = torch.tensor(warp_noise)[..., :3] * torch.tensor(warped_masks)269 warp_noise_vis = (warp_noise_vis - warp_noise_vis.min()) / (warp_noise_vis.max() - warp_noise_vis.min())270 warp_noise_vis = (warp_noise_vis * 255.0).to(torch.uint8)271 272 torchvision.io.write_video(f'{cur_OUTPUT_PATH}/warp_noise.mp4', warp_noise_vis, fps=10, video_codec='h264', options={'crf': '10'})273 274 275 warp_latents = torch.tensor(warp_noise).permute(0, 3, 1, 2).to(latents_org.device).to(latents_org.dtype) # [frame, 4, H, W]276 warp_latents = warp_latents.unsqueeze(0) # [1, frame, 4, H, W]277 278 warped_masks = torch.tensor(warped_masks).permute(0, 3, 1, 2).unsqueeze(0) # [1, frame, 3, H, W]279 mask_extend = torch.concat([warped_masks, warped_masks[:,:,0:1]], dim=2) # [1, frame, 4, H, W]280 mask_extend = mask_extend.to(latents_org.device).to(latents_org.dtype)281 282 warp_latents = warp_latents * mask_extend + latents_org * (1 - mask_extend)283 warp_latents = warp_latents.permute(0, 2, 1, 3, 4)284 random_noise = latents_org.clone().permute(0, 2, 1, 3, 4)285 286 filter_shape = warp_latents.shape287 288 freq_filter = get_freq_filter(289 filter_shape, 290 device = device, 291 filter_type='butterworth',292 n=4,293 d_s=ds,294 d_t=dt295 )296 297 warp_latents = freq_mix_3d(warp_latents, random_noise, freq_filter)298 warp_latents = warp_latents.permute(0, 2, 1, 3, 4)299 300 else:301 warp_latents = latents_org.clone()302 303 generator.manual_seed(42)304 305 with torch.no_grad():306 result = pipeline(307 image=condition_image,308 pose_embedding=cur_plucker_embedding,309 height=height,310 width=width,311 num_frames=num_frames,312 num_inference_steps=num_inference_steps,313 min_guidance_scale=min_guidance_scale,314 max_guidance_scale=max_guidance_scale,315 do_image_process=True,316 generator=generator,317 output_type='pt',318 pose_features= pose_features,319 latents = warp_latents320 ).frames[0].cpu() #[f, c, h, w]321 322 323 result = rearrange(result, 'f c h w -> f h w c')324 result = (result * 255.0).to(torch.uint8)325 326 video_path = tempfile.NamedTemporaryFile(suffix='.mp4').name327 torchvision.io.write_video(video_path, result, fps=10, video_codec='h264', options={'crf': '8'})328 329 return video_path330 331 return run_objctrl_2_5d332 333