finetune/modules/depth_warping/depth_warping.py

from transformers import pipeline
from PIL import Image
import requests
import torchvision
import os
from .camera.WarperPytorch import Warper
import numpy as np
from einops import rearrange, repeat
import torch
import torch.nn as nn
import torch.nn.functional as F
from .depth_anything_v2.dpt import DepthAnythingV2

import pdb

def to_pil_image(x):
    # x: c h w, [-1, 1]
    x_np = ((x+1)/2*255).permute(1,2,0).detach().cpu().numpy().astype(np.uint8)
    x_pil = Image.fromarray(x_np)
    
    return x_pil

def to_npy(x):
    return ((x+1)/2*255).permute(1,2,0).detach().cpu().numpy()

def unnormalize_intrinsic(x, size):
    h, w = size
    x_ = x.detach().clone()
    x_[:,0:1] = x[:,0:1].detach().clone() * w
    x_[:,1:2] = x[:,1:2].detach().clone() * h
    return x_

class DepthWarping_wrapper(nn.Module):
    def __init__(self,
                 model_config,
                 ckpt_path,):
        super().__init__()
        
        # self.depth_model = pipeline(task="depth-estimation", model="depth-anything/Depth-Anything-V2-Small-hf")
        self.depth_model = DepthAnythingV2(**model_config)
        self.depth_model.load_state_dict(torch.load(ckpt_path, map_location='cpu'))
        self.depth_model = self.depth_model.eval()
        self.warper = Warper()
    
    def get_input(self, batch):
        # pdb.set_trace()
        
        b, v = batch["target"]["intrinsics"].shape[:2]
        h, w = batch["context"]["image"].shape[-2:]
        
        
        image = (batch["context"]["image"]) * 2 - 1
        image_ctxt = repeat(image, "b c h w -> (b v) c h w", v=v)
        c2w_ctxt = repeat(batch["context"]["extrinsics"], "b t h w -> (b v t) h w", v=v) # No need to apply inverse as it is an eye matrix.
        # c2w_trgt = rearrange(torch.inverse(batch["target"]["extrinsics"]), "b t h w -> (b t) h w")
        c2w_trgt = rearrange(batch["target"]["extrinsics"], "b t h w -> (b t) h w")
        intrinsics_ctxt = unnormalize_intrinsic(repeat(batch["context"]["intrinsics"], "b t h w -> (b v t) h w", v=v), size=(h,w))
        intrinsics_trgt = unnormalize_intrinsic(rearrange(batch["target"]["intrinsics"], "b t h w -> (b t) h w"), size=(h,w))
        
        # image = image.squeeze(1)
        # depth_ctxt = torch.stack([torch.tensor(self.depth_model.infer_image(to_npy(x))) for x in image], dim=0).to(image.device).unsqueeze(1) # B 1 H W
        depth_ctxt = torch.stack([self.depth_model.infer_image(to_npy(x)) for x in image], dim=0).to(image.device).unsqueeze(1) # B 1 H W

        # depth_ctxt = torch.nn.functional.interpolate(
        #     depth_ctxt,
        #     size=(h,w),
        #     mode="bicubic",
        #     align_corners=False,
        # )
        
        return image_ctxt, c2w_ctxt, c2w_trgt, intrinsics_ctxt, intrinsics_trgt, depth_ctxt, batch['variable_intrinsic']
    
    def forward(self, batch):
        image_ctxt, c2w_ctxt, c2w_trgt, intrinsics_ctxt, intrinsics_trgt, depth_ctxt, variable_intrinsic = self.get_input(batch)
        
        with torch.cuda.amp.autocast(enabled=False):
        
            b, v = batch["target"]["intrinsics"].shape[:2]
            # h, w = image_ctxt.shape[-2:]
            
            warped_trgt, mask_trgt, warped_depth_trgt, flow_f = self.warper.forward_warp(
                frame1=image_ctxt, 
                mask1=None, 
                depth1=repeat(depth_ctxt, "b c h w -> (b t) c h w", t=v), 
                transformation1=c2w_ctxt, 
                transformation2=c2w_trgt, 
                intrinsic1=intrinsics_ctxt, 
                intrinsic2=intrinsics_trgt if variable_intrinsic else None)
            
            warped_src, mask_src, warped_depth_src, flow_b = self.warper.forward_warp(
                frame1=warped_trgt, 
                mask1=None, 
                depth1=warped_depth_trgt, 
                transformation1=c2w_trgt, 
                transformation2=c2w_ctxt, 
                intrinsic1=intrinsics_trgt, 
                intrinsic2=None)
        
        # if use_backward_flow:
        #     mask = mask_trgt
        # else:
        #     mask = mask_src

        return flow_f, flow_b, warped_trgt, depth_ctxt, warped_depth_trgt