import numpy as np import cv2 import PIL from PIL import Image import os from datetime import datetime import pdb import torch.nn.functional as F import numpy as np import os import cv2 import copy from scipy.interpolate import UnivariateSpline, interp1d import numpy as np import PIL.Image import torch import torchvision from tqdm import tqdm from pathlib import Path from typing import Tuple, Optional import cv2 import PIL import numpy import skimage.io import torch import torch.nn.functional as F from decord import VideoReader, cpu from PIL import Image import torch import torch.nn.functional as F # def resize_frames(frames, resolution, device=None): # """ # Resize a batch of frames to (H, W). # Args: # frames: np.ndarray or torch.Tensor # - [T, H, W, C] (channels-last, common for numpy/video) # - [T, C, H, W] (channels-first, PyTorch standard) # - [H, W, C] (single frame, channels-last) # - [C, H, W] (single frame, channels-first) # resolution: tuple (H, W) # device: "cpu", "cuda", or None (auto-detect) # """ # # Convert numpy → torch # if not torch.is_tensor(frames): # frames = torch.from_numpy(frames.astype("float32")) # # Add batch dim if missing # if frames.ndim == 3: # # Single frame [H, W, C] or [C, H, W] → [1, H, W, C] or [1, C, H, W] # frames = frames.unsqueeze(0) # # Ensure channels-first [T, C, H, W] # print("frames shape", frames.shape) # # if frames.shape[-1] in [1, 3]: # [T, H, W, C] # # frames = frames.permute(0, 3, 1, 2) # print("frames shape", frames.shape) # # Normalize dtype # frames = frames.to(torch.float32) # # Send to device # if device is None: # device = "cuda" if torch.cuda.is_available() else "cpu" # frames = frames.to(device) # # Resize # resized = torch.stack([ # torch.from_numpy( # cv2.resize(f.cpu().numpy(), (resolution[1], resolution[0]), interpolation=cv2.INTER_AREA) # ) # for f in frames # ]).to(frames.device).float() # # frames in [T, H, W, C] # # resized = resized.permute(0, 3, 1, 2) # [T, C, H, W] # print("resized shape", resized.shape) # return resized def resize_frames(frames, resolution): Transoposed = False if not torch.is_tensor(frames): print(" frames are not tensors but numpy arrays ") if frames.shape[1]==1 or frames.shape[1]==3: # [T, C, H, W], convert to [T, H, W, C] frames = frames.transpose(0, 2, 3, 1) Transoposed = True resized_frames = [] for frame in frames: img = Image.fromarray((frame * 255).astype('uint8')) img = img.resize((resolution[1], resolution[0]), Image.BILINEAR) img = np.array(img).astype("float32") / 255.0 resized_frames.append(img) resized_frames = np.stack(resized_frames, axis=0) if Transoposed: resized_frames = resized_frames.transpose(0, 3, 1, 2) else: if frames.shape[-1]==1 or frames.shape[-1]==3: # [T, C, H, W], convert to [T, H, W, C] frames = frames.permute(0, 3, 2, 1) Transoposed = True resized_frames = F.interpolate(frames, size=resolution, mode='bilinear', align_corners=True) if Transoposed: resized_frames = resized_frames.permute(0, 2, 3, 1) return resized_frames def read_video_frames_av(video_path, process_length, stride, sample_size, max_res, dataset="open"): # def is_image(path): # return any(path.lower().endswith(ext) for ext in ['.jpg', '.jpeg', '.png', '.bmp']) div = 16 # if is_image(video_path): # print("==> Detected image. Loading as single-frame video:", video_path) # img = Image.open(video_path).convert("RGB") # original_width = img.width # original_height = img.height # height = round(original_height / div) * div # width = round(original_width / div) * div # if max(height, width) > max_res: # scale = max_res / max(original_height, original_width) # height = round(original_height * scale / div) * div # width = round(original_width * scale / div) * div # img = img.resize((width, height), Image.BICUBIC) # img = np.array(img).astype("float32") / 255.0 # [H, W, 3] # frames = img[None, ...] # [1, H, W, 3] # print(f"==> image shape: {frames.shape}") # return frames if dataset == "open": print("==> processing video:", video_path) vid = VideoReader(video_path, ctx=cpu(0)) print("==> original video shape:", (len(vid), *vid.get_batch([0]).shape[1:])) original_height, original_width = vid.get_batch([0]).shape[1:3] # 640, 1080 original_height, original_width = sample_size[0], sample_size[1] height = int(original_height / div) * div # == 640 width = int(original_width / div) * div # == 1088 if max(height, width) > max_res: scale = max_res / max(original_height, original_width) # scale = max_res / max(height, width) height = round(original_height * scale / div) * div width = round(original_width * scale / div) * div vid = VideoReader(video_path, ctx=cpu(0), width=width, height=height) frames_idx = list(range(0, len(vid), stride)) print(f"==> downsampled shape: {(len(frames_idx), *vid.get_batch([0]).shape[1:])}, with stride: {stride}") if process_length != -1 and process_length < len(frames_idx): frames_idx = frames_idx[:process_length] print(f"==> final processing shape: {(len(frames_idx), *vid.get_batch([0]).shape[1:])}") frames = vid.get_batch(frames_idx).asnumpy().astype("float32") / 255.0 # [T, H, W, 3] return frames def read_video_frames(video_path, process_length, stride, max_res, dataset="open"): def is_image(path): return any(path.lower().endswith(ext) for ext in ['.jpg', '.jpeg', '.png', '.bmp']) if is_image(video_path): print("==> Detected image. Loading as single-frame video:", video_path) img = Image.open(video_path).convert("RGB") original_width = img.width original_height = img.height height = round(original_height / 64) * 64 width = round(original_width / 64) * 64 if max(height, width) > max_res: scale = max_res / max(original_height, original_width) height = round(original_height * scale / 64) * 64 width = round(original_width * scale / 64) * 64 img = img.resize((width, height), Image.BICUBIC) img = np.array(img).astype("float32") / 255.0 # [H, W, 3] frames = img[None, ...] # [1, H, W, 3] print(f"==> image shape: {frames.shape}") return frames if dataset == "open": print("==> processing video:", video_path) vid = VideoReader(video_path, ctx=cpu(0)) print("==> original video shape:", (len(vid), *vid.get_batch([0]).shape[1:])) original_height, original_width = vid.get_batch([0]).shape[1:3] height = round(original_height / 64) * 64 width = round(original_width / 64) * 64 if max(height, width) > max_res: scale = max_res / max(original_height, original_width) height = round(original_height * scale / 64) * 64 width = round(original_width * scale / 64) * 64 vid = VideoReader(video_path, ctx=cpu(0), width=width, height=height) frames_idx = list(range(0, len(vid), stride)) print(f"==> downsampled shape: {(len(frames_idx), *vid.get_batch([0]).shape[1:])}, with stride: {stride}") if process_length != -1 and process_length < len(frames_idx): frames_idx = frames_idx[:process_length] print(f"==> final processing shape: {(len(frames_idx), *vid.get_batch([0]).shape[1:])}") frames = vid.get_batch(frames_idx).asnumpy().astype("float32") / 255.0 # [T, H, W, 3] return frames # def read_video_frames(video_path, process_length, stride, max_res, dataset="open"): # def is_image(path): # return any(path.lower().endswith(ext) for ext in ['.jpg', '.jpeg', '.png', '.bmp']) # if is_image(video_path): # print("==> Detected image. Loading as single-frame video:", video_path) # img = Image.open(video_path).convert("RGB") # # FIXME: hard coded # width = 1024 # height = 576 # img = img.resize((width, height), Image.BICUBIC) # img = np.array(img).astype("float32") / 255.0 # [H, W, 3] # frames = img[None, ...] # [1, H, W, 3] # print(f"==> image shape: {frames.shape}") # return frames # if dataset == "open": # print("==> processing video:", video_path) # vid = VideoReader(video_path, ctx=cpu(0)) # print("==> original video shape:", (len(vid), *vid.get_batch([0]).shape[1:])) # # FIXME: hard coded # width = 1024 # height = 576 # vid = VideoReader(video_path, ctx=cpu(0), width=width, height=height) # frames_idx = list(range(0, len(vid), stride)) # print(f"==> downsampled shape: {(len(frames_idx), *vid.get_batch([0]).shape[1:])}, with stride: {stride}") # if process_length != -1 and process_length < len(frames_idx): # frames_idx = frames_idx[:process_length] # print(f"==> final processing shape: {(len(frames_idx), *vid.get_batch([0]).shape[1:])}") # frames = vid.get_batch(frames_idx).asnumpy().astype("float32") / 255.0 # [T, H, W, 3] # return frames def save_video(data, images_path, folder=None, fps=8): if isinstance(data, np.ndarray): tensor_data = (torch.from_numpy(data) * 255).to(torch.uint8) elif isinstance(data, torch.Tensor): tensor_data = (data.detach().cpu() * 255).to(torch.uint8) elif isinstance(data, list): folder = [folder] * len(data) images = [ np.array(Image.open(os.path.join(folder_name, path))) for folder_name, path in zip(folder, data) ] stacked_images = np.stack(images, axis=0) tensor_data = torch.from_numpy(stacked_images).to(torch.uint8) torchvision.io.write_video( images_path, tensor_data, fps=fps, video_codec='h264', options={'crf': '10'} ) def sphere2pose(c2ws_input, theta, phi, r, device, x=None, y=None): c2ws = copy.deepcopy(c2ws_input) # c2ws[:,2, 3] = c2ws[:,2, 3] - radius # 先沿着世界坐标系z轴方向平移再旋转 c2ws[:, 2, 3] -= r if x is not None: c2ws[:, 1, 3] += y if y is not None: c2ws[:, 0, 3] -= x theta = torch.deg2rad(torch.tensor(theta)).to(device) sin_value_x = torch.sin(theta) cos_value_x = torch.cos(theta) rot_mat_x = ( torch.tensor( [ [1, 0, 0, 0], [0, cos_value_x, -sin_value_x, 0], [0, sin_value_x, cos_value_x, 0], [0, 0, 0, 1], ] ) .unsqueeze(0) .repeat(c2ws.shape[0], 1, 1) .to(device) ) phi = torch.deg2rad(torch.tensor(phi)).to(device) sin_value_y = torch.sin(phi) cos_value_y = torch.cos(phi) rot_mat_y = ( torch.tensor( [ [cos_value_y, 0, sin_value_y, 0], [0, 1, 0, 0], [-sin_value_y, 0, cos_value_y, 0], [0, 0, 0, 1], ] ) .unsqueeze(0) .repeat(c2ws.shape[0], 1, 1) .to(device) ) c2ws = torch.matmul(rot_mat_x, c2ws) c2ws = torch.matmul(rot_mat_y, c2ws) # c2ws[:,2, 3] = c2ws[:,2, 3] + radius return c2ws def generate_traj_specified(c2ws_anchor, theta, phi, d_r, d_x, d_y, frame, device, opts): # Initialize a camera. camera_x = getattr(opts, "init_dx", 0.0) camera_y = getattr(opts, "init_dy", 0.0) camera_z = getattr(opts, "init_dz", 0.0) camera_theta = getattr(opts, "init_theta", 0.0) camera_phi = getattr(opts, "init_phi", 0.0) thetas = np.linspace(0, theta, frame) + camera_theta phis = np.linspace(0, phi, frame) + camera_phi rs = np.linspace(0, d_r, frame)+camera_z xs = np.linspace(0, d_x, frame)+camera_x ys = np.linspace(0, d_y, frame)+camera_y c2ws_list = [] c2ws_list.append(copy.deepcopy(c2ws_anchor)) for th, ph, r, x, y in zip(thetas, phis, rs, xs, ys): c2w_new = sphere2pose( c2ws_anchor, np.float32(th), np.float32(ph), np.float32(r), device, np.float32(x), np.float32(y), ) c2ws_list.append(c2w_new) c2ws = torch.cat(c2ws_list, dim=0) return c2ws def generate_traj_specified_frame(c2ws_anchor, theta, phi, d_r, d_x, d_y, frame, device, opts): # Initialize a camera. camera_x = getattr(opts, "init_dx", 0.0) camera_y = getattr(opts, "init_dy", 0.0) camera_z = getattr(opts, "init_dz", 0.0) camera_theta = getattr(opts, "init_theta", 0.0) camera_phi = getattr(opts, "init_phi", 0.0) # thetas = np.linspace(0, theta, frame) + camera_theta # phis = np.linspace(0, phi, frame) + camera_phi # rs = np.linspace(0, d_r, frame)+camera_z # xs = np.linspace(0, d_x, frame)+camera_x # ys = np.linspace(0, d_y, frame)+camera_y thetas = np.full(frame, theta/(frame)) + camera_theta phis = np.full(frame, phi/(frame)) + camera_phi rs = np.full(frame, d_r/(frame)) + camera_z xs = np.full(frame, d_x/(frame)) + camera_x ys = np.full(frame, d_y/(frame)) + camera_y c2ws_list = [] c2ws_list.append(copy.deepcopy(c2ws_anchor)) for th, ph, r, x, y in zip(thetas, phis, rs, xs, ys): c2w_new = sphere2pose( c2ws_anchor, np.float32(th), np.float32(ph), np.float32(r), device, np.float32(x), np.float32(y), ) c2ws_list.append(c2w_new) c2ws = torch.cat(c2ws_list, dim=0) return c2ws def generate_traj_specified_fast(c2ws_anchor, theta, phi, d_r, d_x, d_y, frame, device): half = frame // 3 thetas = np.linspace(0, theta, half) phis = np.linspace(0, phi, half) rs = np.linspace(0, d_r, half) xs = np.linspace(0, d_x, half) ys = np.linspace(0, d_y, half) c2ws_list = [] for th, ph, r, x, y in zip(thetas, phis, rs, xs, ys): c2w_new = sphere2pose( c2ws_anchor, np.float32(th), np.float32(ph), np.float32(r), device, np.float32(x), np.float32(y), ) c2ws_list.append(c2w_new) last_c2w = c2ws_list[-1] for _ in range(frame - half): c2ws_list.append(last_c2w.clone()) c2ws = torch.cat(c2ws_list, dim=0) return c2ws def txt_interpolation(input_list, n, mode='smooth'): x = np.linspace(0, 1, len(input_list)) if mode == 'smooth': f = UnivariateSpline(x, input_list, k=3) elif mode == 'linear': f = interp1d(x, input_list) else: raise KeyError(f"Invalid txt interpolation mode: {mode}") xnew = np.linspace(0, 1, n) ynew = f(xnew) return ynew def generate_traj_txt(c2ws_anchor, phi, theta, r, frame, device, x = None, y = None): # Initialize a camera. """ The camera coordinate sysmte in COLMAP is right-down-forward Pytorch3D is left-up-forward """ if len(phi) > 3: phis = txt_interpolation(phi, frame, mode='smooth') phis[0] = phi[0] phis[-1] = phi[-1] else: phis = txt_interpolation(phi, frame, mode='linear') if len(theta) > 3: thetas = txt_interpolation(theta, frame, mode='smooth') thetas[0] = theta[0] thetas[-1] = theta[-1] else: thetas = txt_interpolation(theta, frame, mode='linear') if len(r) > 3: rs = txt_interpolation(r, frame, mode='smooth') rs[0] = r[0] rs[-1] = r[-1] else: rs = txt_interpolation(r, frame, mode='linear') if len(x) > 3: xs = txt_interpolation(x, frame, mode='smooth') xs[0] = x[0] xs[-1] = x[-1] else: xs = txt_interpolation(x, frame, mode='linear') if len(y) > 3: ys = txt_interpolation(y, frame, mode='smooth') ys[0] = y[0] ys[-1] = y[-1] else: ys = txt_interpolation(y, frame, mode='linear') # rs = rs*c2ws_anchor[0,2,3].cpu().numpy() c2ws_list = [] c2ws_list.append(copy.deepcopy(c2ws_anchor)) for th, ph, r, x_i, y_i in zip(thetas, phis, rs, xs, ys): c2w_new = sphere2pose( c2ws_anchor, np.float32(th), np.float32(ph), np.float32(r), device, np.float32(x_i), np.float32(y_i) ) c2ws_list.append(c2w_new) c2ws = torch.cat(c2ws_list, dim=0) return c2ws class Warper: def __init__(self, resolution: tuple = None, device: str = 'gpu0'): self.resolution = resolution self.device = self.get_device(device) self.dtype = torch.float32 return def forward_warp( self, frame1: torch.Tensor, mask1: Optional[torch.Tensor], depth1: torch.Tensor, transformation1: torch.Tensor, transformation2: torch.Tensor, intrinsic1: torch.Tensor, intrinsic2: Optional[torch.Tensor], mask=False, twice=False ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """ Given a frame1 and global transformations transformation1 and transformation2, warps frame1 to next view using bilinear splatting. All arrays should be torch tensors with batch dimension and channel first :param frame1: (b, 3, h, w). If frame1 is not in the range [-1, 1], either set is_image=False when calling bilinear_splatting on frame within this function, or modify clipping in bilinear_splatting() method accordingly. :param mask1: (b, 1, h, w) - 1 for known, 0 for unknown. Optional :param depth1: (b, 1, h, w) :param transformation1: (b, 4, 4) extrinsic transformation matrix of first view: [R, t; 0, 1] :param transformation2: (b, 4, 4) extrinsic transformation matrix of second view: [R, t; 0, 1] :param intrinsic1: (b, 3, 3) camera intrinsic matrix :param intrinsic2: (b, 3, 3) camera intrinsic matrix. Optional """ if self.resolution is not None: assert frame1.shape[2:4] == self.resolution b, c, h, w = frame1.shape if mask1 is None: mask1 = torch.ones(size=(b, 1, h, w)).to(frame1) if intrinsic2 is None: intrinsic2 = intrinsic1.clone() assert frame1.shape == (b, 3, h, w) assert mask1.shape == (b, 1, h, w) assert depth1.shape == (b, 1, h, w) assert transformation1.shape == (b, 4, 4) assert transformation2.shape == (b, 4, 4) assert intrinsic1.shape == (b, 3, 3) assert intrinsic2.shape == (b, 3, 3) frame1 = frame1.to(self.device).to(self.dtype) mask1 = mask1.to(self.device).to(self.dtype) depth1 = depth1.to(self.device).to(self.dtype) transformation1 = transformation1.to(self.device).to(self.dtype) transformation2 = transformation2.to(self.device).to(self.dtype) intrinsic1 = intrinsic1.to(self.device).to(self.dtype) intrinsic2 = intrinsic2.to(self.device).to(self.dtype) # need to change this function as well. trans_points1 = self.compute_transformed_points( depth1, transformation1, transformation2, intrinsic1, intrinsic2 ) trans_coordinates = ( trans_points1[:, :, :, :2, 0] / trans_points1[:, :, :, 2:3, 0] ) trans_depth1 = trans_points1[:, :, :, 2, 0] grid = self.create_grid(b, h, w).to(trans_coordinates) flow12 = trans_coordinates.permute(0, 3, 1, 2) - grid if not twice: warped_frame2, mask2 = self.bilinear_splatting( frame1, mask1, trans_depth1, flow12, None, is_image=True ) if mask: warped_frame2, mask2 = self.clean_points(warped_frame2, mask2) return warped_frame2, mask2, None, flow12 else: warped_frame2, mask2 = self.bilinear_splatting( frame1, mask1, trans_depth1, flow12, None, is_image=True ) warped_flow, _ = self.bilinear_splatting( flow12, mask1, trans_depth1, flow12, None, is_image=False ) twice_warped_frame1, _ = self.bilinear_splatting( warped_frame2, mask2, depth1.squeeze(1), -warped_flow, None, is_image=True, ) return twice_warped_frame1, warped_frame2, None, None def compute_transformed_points( self, depth1: torch.Tensor, transformation1: torch.Tensor, transformation2: torch.Tensor, intrinsic1: torch.Tensor, intrinsic2: Optional[torch.Tensor], ): """ Computes transformed position for each pixel location """ if self.resolution is not None: assert depth1.shape[2:4] == self.resolution b, _, h, w = depth1.shape if intrinsic2 is None: intrinsic2 = intrinsic1.clone() transformation = torch.bmm( transformation2, torch.linalg.inv(transformation1) ) # (b, 4, 4) x1d = torch.arange(0, w)[None] y1d = torch.arange(0, h)[:, None] x2d = x1d.repeat([h, 1]).to(depth1) # (h, w) y2d = y1d.repeat([1, w]).to(depth1) # (h, w) ones_2d = torch.ones(size=(h, w)).to(depth1) # (h, w) ones_4d = ones_2d[None, :, :, None, None].repeat( [b, 1, 1, 1, 1] ) # (b, h, w, 1, 1) pos_vectors_homo = torch.stack([x2d, y2d, ones_2d], dim=2)[ None, :, :, :, None ] # (1, h, w, 3, 1) intrinsic1_inv = torch.linalg.inv(intrinsic1) # (b, 3, 3) intrinsic1_inv_4d = intrinsic1_inv[:, None, None] # (b, 1, 1, 3, 3) intrinsic2_4d = intrinsic2[:, None, None] # (b, 1, 1, 3, 3) depth_4d = depth1[:, 0][:, :, :, None, None] # (b, h, w, 1, 1) trans_4d = transformation[:, None, None] # (b, 1, 1, 4, 4) unnormalized_pos = torch.matmul( intrinsic1_inv_4d, pos_vectors_homo ) # (b, h, w, 3, 1) world_points = depth_4d * unnormalized_pos # (b, h, w, 3, 1) world_points_homo = torch.cat([world_points, ones_4d], dim=3) # (b, h, w, 4, 1) trans_world_homo = torch.matmul(trans_4d, world_points_homo) # (b, h, w, 4, 1) trans_world = trans_world_homo[:, :, :, :3] # (b, h, w, 3, 1) trans_norm_points = torch.matmul(intrinsic2_4d, trans_world) # (b, h, w, 3, 1) return trans_norm_points def compute_transformed_points_frame( self, depth1: torch.Tensor, transformation1: torch.Tensor, transformation2: torch.Tensor, intrinsic1: torch.Tensor, intrinsic2: Optional[torch.Tensor], grid: Optional[torch.Tensor] = None, # NEW: grid (2, h, w), grid[0]=x, grid[1]=y ): """ Computes transformed positions for each pixel location or provided grid. Args: depth1: (b, 1, h, w) depth map transformation1: (b, 4, 4) camera-to-world transform of source transformation2: (b, 4, 4) camera-to-world transform of target intrinsic1: (b, 3, 3) intrinsics of source intrinsic2: (b, 3, 3) intrinsics of target (optional) grid: (2, h, w), where grid[0] = x coords, grid[1] = y coords """ if self.resolution is not None: assert depth1.shape[2:4] == self.resolution b, _, h, w = depth1.shape if intrinsic2 is None: intrinsic2 = intrinsic1.clone() # relative transformation transformation = torch.bmm( transformation2, torch.linalg.inv(transformation1) ) # (b, 4, 4) # If grid is not provided, build pixel grid if grid is None: x1d = torch.arange(0, w, device=depth1.device)[None] y1d = torch.arange(0, h, device=depth1.device)[:, None] x2d = x1d.repeat([h, 1]) # (h, w) y2d = y1d.repeat([1, w]) # (h, w) else: x2d = grid[0][0].to(depth1.device) # (h, w) y2d = grid[0][1].to(depth1.device) # (h, w) # build homogeneous coordinates from grid ones_2d = torch.ones(size=(h, w), device=depth1.device) pos_vectors_homo = torch.stack([x2d, y2d, ones_2d], dim=2)[ None, :, :, :, None ] # (1, h, w, 3, 1) # expand for batch ones_4d = ones_2d[None, :, :, None, None].repeat([b, 1, 1, 1, 1]) # (b,h,w,1,1) intrinsic1_inv = torch.linalg.inv(intrinsic1) # (b,3,3) intrinsic1_inv_4d = intrinsic1_inv[:, None, None] # (b,1,1,3,3) intrinsic2_4d = intrinsic2[:, None, None] # (b,1,1,3,3) depth_4d = depth1[:, 0][:, :, :, None, None] # (b,h,w,1,1) trans_4d = transformation[:, None, None] # (b,1,1,4,4) # backproject unnormalized_pos = torch.matmul(intrinsic1_inv_4d, pos_vectors_homo) # (b,h,w,3,1) world_points = depth_4d * unnormalized_pos # (b,h,w,3,1) world_points_homo = torch.cat([world_points, ones_4d], dim=3) # (b,h,w,4,1) # transform & reproject trans_world_homo = torch.matmul(trans_4d, world_points_homo) # (b,h,w,4,1) trans_world = trans_world_homo[:, :, :, :3] # (b,h,w,3,1) trans_norm_points = torch.matmul(intrinsic2_4d, trans_world) # (b,h,w,3,1) return trans_norm_points def bilinear_splatting( self, frame1: torch.Tensor, mask1: Optional[torch.Tensor], depth1: torch.Tensor, flow12: torch.Tensor, flow12_mask: Optional[torch.Tensor], is_image: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Bilinear splatting :param frame1: (b,c,h,w) :param mask1: (b,1,h,w): 1 for known, 0 for unknown. Optional :param depth1: (b,1,h,w) :param flow12: (b,2,h,w) :param flow12_mask: (b,1,h,w): 1 for valid flow, 0 for invalid flow. Optional :param is_image: if true, output will be clipped to (-1,1) range :return: warped_frame2: (b,c,h,w) mask2: (b,1,h,w): 1 for known and 0 for unknown """ if self.resolution is not None: assert frame1.shape[2:4] == self.resolution b, c, h, w = frame1.shape if mask1 is None: mask1 = torch.ones(size=(b, 1, h, w)).to(frame1) if flow12_mask is None: flow12_mask = torch.ones(size=(b, 1, h, w)).to(flow12) grid = self.create_grid(b, h, w).to(frame1) trans_pos = flow12 + grid trans_pos_offset = trans_pos + 1 trans_pos_floor = torch.floor(trans_pos_offset).long() trans_pos_ceil = torch.ceil(trans_pos_offset).long() trans_pos_offset = torch.stack( [ torch.clamp(trans_pos_offset[:, 0], min=0, max=w + 1), torch.clamp(trans_pos_offset[:, 1], min=0, max=h + 1), ], dim=1, ) trans_pos_floor = torch.stack( [ torch.clamp(trans_pos_floor[:, 0], min=0, max=w + 1), torch.clamp(trans_pos_floor[:, 1], min=0, max=h + 1), ], dim=1, ) trans_pos_ceil = torch.stack( [ torch.clamp(trans_pos_ceil[:, 0], min=0, max=w + 1), torch.clamp(trans_pos_ceil[:, 1], min=0, max=h + 1), ], dim=1, ) prox_weight_nw = (1 - (trans_pos_offset[:, 1:2] - trans_pos_floor[:, 1:2])) * ( 1 - (trans_pos_offset[:, 0:1] - trans_pos_floor[:, 0:1]) ) prox_weight_sw = (1 - (trans_pos_ceil[:, 1:2] - trans_pos_offset[:, 1:2])) * ( 1 - (trans_pos_offset[:, 0:1] - trans_pos_floor[:, 0:1]) ) prox_weight_ne = (1 - (trans_pos_offset[:, 1:2] - trans_pos_floor[:, 1:2])) * ( 1 - (trans_pos_ceil[:, 0:1] - trans_pos_offset[:, 0:1]) ) prox_weight_se = (1 - (trans_pos_ceil[:, 1:2] - trans_pos_offset[:, 1:2])) * ( 1 - (trans_pos_ceil[:, 0:1] - trans_pos_offset[:, 0:1]) ) sat_depth1 = torch.clamp(depth1, min=0, max=1000) log_depth1 = torch.log(1 + sat_depth1) depth_weights = torch.exp(log_depth1 / log_depth1.max() * 50) weight_nw = torch.moveaxis( prox_weight_nw * mask1 * flow12_mask / depth_weights.unsqueeze(1), [0, 1, 2, 3], [0, 3, 1, 2], ) weight_sw = torch.moveaxis( prox_weight_sw * mask1 * flow12_mask / depth_weights.unsqueeze(1), [0, 1, 2, 3], [0, 3, 1, 2], ) weight_ne = torch.moveaxis( prox_weight_ne * mask1 * flow12_mask / depth_weights.unsqueeze(1), [0, 1, 2, 3], [0, 3, 1, 2], ) weight_se = torch.moveaxis( prox_weight_se * mask1 * flow12_mask / depth_weights.unsqueeze(1), [0, 1, 2, 3], [0, 3, 1, 2], ) warped_frame = torch.zeros(size=(b, h + 2, w + 2, c), dtype=torch.float32).to( frame1 ) warped_weights = torch.zeros(size=(b, h + 2, w + 2, 1), dtype=torch.float32).to( frame1 ) frame1_cl = torch.moveaxis(frame1, [0, 1, 2, 3], [0, 3, 1, 2]) batch_indices = torch.arange(b)[:, None, None].to(frame1.device) warped_frame.index_put_( (batch_indices, trans_pos_floor[:, 1], trans_pos_floor[:, 0]), frame1_cl * weight_nw, accumulate=True, ) warped_frame.index_put_( (batch_indices, trans_pos_ceil[:, 1], trans_pos_floor[:, 0]), frame1_cl * weight_sw, accumulate=True, ) warped_frame.index_put_( (batch_indices, trans_pos_floor[:, 1], trans_pos_ceil[:, 0]), frame1_cl * weight_ne, accumulate=True, ) warped_frame.index_put_( (batch_indices, trans_pos_ceil[:, 1], trans_pos_ceil[:, 0]), frame1_cl * weight_se, accumulate=True, ) warped_weights.index_put_( (batch_indices, trans_pos_floor[:, 1], trans_pos_floor[:, 0]), weight_nw, accumulate=True, ) warped_weights.index_put_( (batch_indices, trans_pos_ceil[:, 1], trans_pos_floor[:, 0]), weight_sw, accumulate=True, ) warped_weights.index_put_( (batch_indices, trans_pos_floor[:, 1], trans_pos_ceil[:, 0]), weight_ne, accumulate=True, ) warped_weights.index_put_( (batch_indices, trans_pos_ceil[:, 1], trans_pos_ceil[:, 0]), weight_se, accumulate=True, ) warped_frame_cf = torch.moveaxis(warped_frame, [0, 1, 2, 3], [0, 2, 3, 1]) warped_weights_cf = torch.moveaxis(warped_weights, [0, 1, 2, 3], [0, 2, 3, 1]) cropped_warped_frame = warped_frame_cf[:, :, 1:-1, 1:-1] cropped_weights = warped_weights_cf[:, :, 1:-1, 1:-1] mask = cropped_weights > 0 zero_value = -1 if is_image else 0 zero_tensor = torch.tensor(zero_value, dtype=frame1.dtype, device=frame1.device) warped_frame2 = torch.where( mask, cropped_warped_frame / cropped_weights, zero_tensor ) mask2 = mask.to(frame1) if is_image: assert warped_frame2.min() >= -1.1 # Allow for rounding errors assert warped_frame2.max() <= 1.1 warped_frame2 = torch.clamp(warped_frame2, min=-1, max=1) return warped_frame2, mask2 def clean_points(self, warped_frame2, mask2): warped_frame2 = (warped_frame2 + 1.0) / 2.0 mask = 1 - mask2 mask[mask < 0.5] = 0 mask[mask >= 0.5] = 1 mask = mask.squeeze(0).repeat(3, 1, 1).permute(1, 2, 0) * 255.0 mask = mask.cpu().numpy() kernel = numpy.ones((5, 5), numpy.uint8) mask_erosion = cv2.dilate(numpy.array(mask), kernel, iterations=1) mask_erosion = PIL.Image.fromarray(numpy.uint8(mask_erosion)) mask_erosion_ = numpy.array(mask_erosion) / 255.0 mask_erosion_[mask_erosion_ < 0.5] = 0 mask_erosion_[mask_erosion_ >= 0.5] = 1 mask_new = ( torch.from_numpy(mask_erosion_) .permute(2, 0, 1) .unsqueeze(0) .to(self.device) ) warped_frame2 = warped_frame2 * (1 - mask_new) return warped_frame2 * 2.0 - 1.0, 1 - mask_new[:, 0:1, :, :] @staticmethod def create_grid(b, h, w): x_1d = torch.arange(0, w)[None] y_1d = torch.arange(0, h)[:, None] x_2d = x_1d.repeat([h, 1]) y_2d = y_1d.repeat([1, w]) grid = torch.stack([x_2d, y_2d], dim=0) batch_grid = grid[None].repeat([b, 1, 1, 1]) return batch_grid @staticmethod def read_image(path: Path) -> torch.Tensor: image = skimage.io.imread(path.as_posix()) return image @staticmethod def read_depth(path: Path) -> torch.Tensor: if path.suffix == '.png': depth = skimage.io.imread(path.as_posix()) elif path.suffix == '.npy': depth = numpy.load(path.as_posix()) elif path.suffix == '.npz': with numpy.load(path.as_posix()) as depth_data: depth = depth_data['depth'] else: raise RuntimeError(f'Unknown depth format: {path.suffix}') return depth @staticmethod def camera_intrinsic_transform( capture_width=1920, capture_height=1080, patch_start_point: tuple = (0, 0) ): start_y, start_x = patch_start_point camera_intrinsics = numpy.eye(4) camera_intrinsics[0, 0] = 2100 camera_intrinsics[0, 2] = capture_width / 2.0 - start_x camera_intrinsics[1, 1] = 2100 camera_intrinsics[1, 2] = capture_height / 2.0 - start_y return camera_intrinsics @staticmethod def get_device(device: str): """ Returns torch device object :param device: cpu/gpu0/gpu1 :return: """ if device == 'cpu': device = torch.device('cpu') elif device.startswith('gpu') and torch.cuda.is_available(): gpu_num = int(device[3:]) device = torch.device(f'cuda:{gpu_num}') else: device = torch.device('cpu') return device