hugging-apps/echo-memory

1import torch2import torch.nn as nn3import numpy as np4from einops import rearrange5import os6from typing_extensions import Literal7 8class SimpleAdapter(nn.Module):9    def __init__(self, in_dim, out_dim, kernel_size, stride, num_residual_blocks=1):10        super(SimpleAdapter, self).__init__()11 12        # Pixel Unshuffle: reduce spatial dimensions by a factor of 813        self.pixel_unshuffle = nn.PixelUnshuffle(downscale_factor=8)14 15        # Convolution: reduce spatial dimensions by a factor16        #  of 2 (without overlap)17        self.conv = nn.Conv2d(in_dim * 64, out_dim, kernel_size=kernel_size, stride=stride, padding=0)18 19        # Residual blocks for feature extraction20        self.residual_blocks = nn.Sequential(21            *[ResidualBlock(out_dim) for _ in range(num_residual_blocks)]22        )23 24    def forward(self, x):25        # Reshape to merge the frame dimension into batch26        bs, c, f, h, w = x.size()27        x = x.permute(0, 2, 1, 3, 4).contiguous().view(bs * f, c, h, w)28 29        # Pixel Unshuffle operation30        x_unshuffled = self.pixel_unshuffle(x)31 32        # Convolution operation33        x_conv = self.conv(x_unshuffled)34 35        # Feature extraction with residual blocks36        out = self.residual_blocks(x_conv)37 38        # Reshape to restore original bf dimension39        out = out.view(bs, f, out.size(1), out.size(2), out.size(3))40 41        # Permute dimensions to reorder (if needed), e.g., swap channels and feature frames42        out = out.permute(0, 2, 1, 3, 4)43 44        return out45    46    def process_camera_coordinates(47        self,48        direction: Literal["Left", "Right", "Up", "Down", "LeftUp", "LeftDown", "RightUp", "RightDown"],49        length: int,50        height: int,51        width: int,52        speed: float = 1/54,53        origin=(0, 0.532139961, 0.946026558, 0.5, 0.5, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0)54    ):55        if origin is None:56            origin = (0, 0.532139961, 0.946026558, 0.5, 0.5, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0)57        coordinates = generate_camera_coordinates(direction, length, speed, origin)58        plucker_embedding = process_pose_file(coordinates, width, height)59        return plucker_embedding60        61    62 63class ResidualBlock(nn.Module):64    def __init__(self, dim):65        super(ResidualBlock, self).__init__()66        self.conv1 = nn.Conv2d(dim, dim, kernel_size=3, padding=1)67        self.relu = nn.ReLU(inplace=True)68        self.conv2 = nn.Conv2d(dim, dim, kernel_size=3, padding=1)69 70    def forward(self, x):71        residual = x72        out = self.relu(self.conv1(x))73        out = self.conv2(out)74        out += residual75        return out76    77class Camera(object):78    """Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py79    """80    def __init__(self, entry):81        fx, fy, cx, cy = entry[1:5]82        self.fx = fx83        self.fy = fy84        self.cx = cx85        self.cy = cy86        w2c_mat = np.array(entry[7:]).reshape(3, 4)87        w2c_mat_4x4 = np.eye(4)88        w2c_mat_4x4[:3, :] = w2c_mat89        self.w2c_mat = w2c_mat_4x490        self.c2w_mat = np.linalg.inv(w2c_mat_4x4)91 92def get_relative_pose(cam_params):93    """Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py94    """95    abs_w2cs = [cam_param.w2c_mat for cam_param in cam_params]96    abs_c2ws = [cam_param.c2w_mat for cam_param in cam_params]97    cam_to_origin = 098    target_cam_c2w = np.array([99        [1, 0, 0, 0],100        [0, 1, 0, -cam_to_origin],101        [0, 0, 1, 0],102        [0, 0, 0, 1]103    ])104    abs2rel = target_cam_c2w @ abs_w2cs[0]105    ret_poses = [target_cam_c2w, ] + [abs2rel @ abs_c2w for abs_c2w in abs_c2ws[1:]]106    ret_poses = np.array(ret_poses, dtype=np.float32)107    return ret_poses108 109def custom_meshgrid(*args):110    # torch>=2.0.0 only111    return torch.meshgrid(*args, indexing='ij')112 113 114def ray_condition(K, c2w, H, W, device):115    """Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py116    """117    # c2w: B, V, 4, 4118    # K: B, V, 4119 120    B = K.shape[0]121 122    j, i = custom_meshgrid(123        torch.linspace(0, H - 1, H, device=device, dtype=c2w.dtype),124        torch.linspace(0, W - 1, W, device=device, dtype=c2w.dtype),125    )126    i = i.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5  # [B, HxW]127    j = j.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5  # [B, HxW]128 129    fx, fy, cx, cy = K.chunk(4, dim=-1)  # B,V, 1130 131    zs = torch.ones_like(i)  # [B, HxW]132    xs = (i - cx) / fx * zs133    ys = (j - cy) / fy * zs134    zs = zs.expand_as(ys)135 136    directions = torch.stack((xs, ys, zs), dim=-1)  # B, V, HW, 3137    directions = directions / directions.norm(dim=-1, keepdim=True)  # B, V, HW, 3138 139    rays_d = directions @ c2w[..., :3, :3].transpose(-1, -2)  # B, V, 3, HW140    rays_o = c2w[..., :3, 3]  # B, V, 3141    rays_o = rays_o[:, :, None].expand_as(rays_d)  # B, V, 3, HW142    # c2w @ dirctions143    rays_dxo = torch.linalg.cross(rays_o, rays_d)144    plucker = torch.cat([rays_dxo, rays_d], dim=-1)145    plucker = plucker.reshape(B, c2w.shape[1], H, W, 6)  # B, V, H, W, 6146    # plucker = plucker.permute(0, 1, 4, 2, 3)147    return plucker148 149 150def process_pose_file(cam_params, width=672, height=384, original_pose_width=1280, original_pose_height=720, device='cpu', return_poses=False):151    if return_poses:152        return cam_params153    else:154        cam_params = [Camera(cam_param) for cam_param in cam_params]155 156        sample_wh_ratio = width / height157        pose_wh_ratio = original_pose_width / original_pose_height  # Assuming placeholder ratios, change as needed158 159        if pose_wh_ratio > sample_wh_ratio:160            resized_ori_w = height * pose_wh_ratio161            for cam_param in cam_params:162                cam_param.fx = resized_ori_w * cam_param.fx / width163        else:164            resized_ori_h = width / pose_wh_ratio165            for cam_param in cam_params:166                cam_param.fy = resized_ori_h * cam_param.fy / height167 168        intrinsic = np.asarray([[cam_param.fx * width,169                                cam_param.fy * height,170                                cam_param.cx * width,171                                cam_param.cy * height]172                                for cam_param in cam_params], dtype=np.float32)173 174        K = torch.as_tensor(intrinsic)[None]  # [1, 1, 4]175        c2ws = get_relative_pose(cam_params)  # Assuming this function is defined elsewhere176        c2ws = torch.as_tensor(c2ws)[None]  # [1, n_frame, 4, 4]177        plucker_embedding = ray_condition(K, c2ws, height, width, device=device)[0].permute(0, 3, 1, 2).contiguous()  # V, 6, H, W178        plucker_embedding = plucker_embedding[None]179        plucker_embedding = rearrange(plucker_embedding, "b f c h w -> b f h w c")[0]180        return plucker_embedding181 182 183 184def generate_camera_coordinates(185    direction: Literal["Left", "Right", "Up", "Down", "LeftUp", "LeftDown", "RightUp", "RightDown"],186    length: int,187    speed: float = 1/54,188    origin=(0, 0.532139961, 0.946026558, 0.5, 0.5, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0)189):190    coordinates = [list(origin)]191    while len(coordinates) < length:192        coor = coordinates[-1].copy()193        if "Left" in direction:194            coor[9] += speed195        if "Right" in direction:196            coor[9] -= speed197        if "Up" in direction:198            coor[13] += speed199        if "Down" in direction:200            coor[13] -= speed201        coordinates.append(coor)202    return coordinates203