lib/midas.py

import os
import glob
import torch
import cv2
import matplotlib.pyplot as plt
import os

import numpy as np
import torch.fft as fft
import ipdb
import copy
import wget

from midas.model_loader import load_model
import torch.nn.functional as F
first_execution = True
thisdir = os.path.abspath(os.path.dirname(__file__))

class MiDas():
    def __init__(self, device, model_type) -> None:
        self.device = device

        torch.backends.cudnn.enabled = True
        torch.backends.cudnn.benchmark = True

        model_weights = os.path.join(thisdir, '..' ,f"./weights/{model_type}.pt")
        if not os.path.exists(model_weights):
            os.makedirs(os.path.dirname(model_weights), exist_ok=True)
            if '384' in model_type:
                wget.download('https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_swin2_large_384.pt', model_weights)
            elif '512' in model_type:
                wget.download('https://github.com/isl-org/MiDaS/releases/download/v3_1/dpt_beit_large_512.pt', model_weights)
            else:
                assert False, 'please select correct depth estimation model.'
        print("Device: %s" % device)
        model, transform, net_w, net_h = load_model(
            device, model_weights, model_type, optimize=False, height=None, square=False
        )
        self.model = model
        self.transform = transform
        self.model_type = model_type
        self.net_w = net_w
        self.net_h = net_h

    def process(
        self, device, model, model_type, image, input_size, target_size, optimize, use_camera
    ):
        """
        Run the inference and interpolate.

        Args:
            device (torch.device): the torch device used
            model: the model used for inference
            model_type: the type of the model
            image: the image fed into the neural network
            input_size: the size (width, height) of the neural network input (for OpenVINO)
            target_size: the size (width, height) the neural network output is interpolated to
            optimize: optimize the model to half-floats on CUDA?
            use_camera: is the camera used?

        Returns:
            the prediction
        """
        global first_execution

        if "openvino" in model_type:
            if first_execution or not use_camera:
                # print(
                #     f"    Input resized to {input_size[0]}x{input_size[1]} before entering the encoder"
                # )
                first_execution = False

            sample = [np.reshape(image, (1, 3, *input_size))]
            prediction = model(sample)[model.output(0)][0]
            prediction = cv2.resize(
                prediction, dsize=target_size, interpolation=cv2.INTER_CUBIC
            )
        else:
            sample = torch.from_numpy(image).to(device).unsqueeze(0)

            if optimize and device == torch.device("cuda"):
                if first_execution:
                    print(
                        "  Optimization to half-floats activated. Use with caution, because models like Swin require\n"
                        "  float precision to work properly and may yield non-finite depth values to some extent for\n"
                        "  half-floats."
                    )
                sample = sample.to(memory_format=torch.channels_last)
                sample = sample.half()

            if first_execution or not use_camera:
                height, width = sample.shape[2:]
                print(f"    Input resized to {width}x{height} before entering the encoder")
                first_execution = False

            prediction = model.forward(sample)
            prediction = (
                torch.nn.functional.interpolate(
                    prediction.unsqueeze(1),
                    size=target_size[::-1],
                    mode="bicubic",
                    align_corners=False,
                )
                .squeeze()
                .cpu()
                .numpy()
            )

        return prediction
    
    def prediction2depth(self, depth):
        bits = 1
        if not np.isfinite(depth).all():
            depth=np.nan_to_num(depth, nan=0.0, posinf=0.0, neginf=0.0)
            print("WARNING: Non-finite depth values present")
        
        depth_min = depth.min()
        depth_max = depth.max()

        max_val = (2**(8*bits))-1

        if depth_max - depth_min > np.finfo("float").eps:
            out = max_val * (depth - depth_min) / (depth_max - depth_min)
        else:
            out = np.zeros(depth.shape, dtype=depth.dtype)
        # out = cv2.applyColorMap(np.uint8(out), cv2.COLORMAP_INFERNO)
        return out
    
    def calc_R(self, theta_z, theta_x, theta_y):
        theta_z, theta_x, theta_y = theta_z/180*np.pi, theta_x/180*np.pi, theta_y/180*np.pi, 
        Rz = np.array([[np.cos(theta_z), np.sin(theta_z), 0],
                    [-np.sin(theta_z), np.cos(theta_z), 0],
                    [0,0,1]])
        Rx = np.array([[1,0,0],
                    [0,np.cos(theta_x), np.sin(theta_x)],
                    [0, -np.sin(theta_x), np.cos(theta_x)]])
        Ry = np.array([[np.cos(theta_y), 0, np.sin(theta_y)],
                    [0,1,0],
                    [-np.sin(theta_y), 0, np.cos(theta_y)]])
        
        R = Rz @ Rx @ Ry
        return R 
    
    def render_new_view(self, img, depth, R, t, K):
        h, w, _ = img.shape
        new_img = np.zeros_like(img)

        for y in range(h):
            for x in range(w):
                # Back-project
                Z = depth[y, x]
                X = (x - K[0, 2]) * Z / K[0, 0]
                Y = (y - K[1, 2]) * Z / K[1, 1]
                point3D = np.array([X, Y, Z, 1])

                # Transform
                point3D_new = R @ point3D[:3] + t
                if point3D_new[2] <= 0:  # point is behind the camera
                    continue

                # Project to new view
                u = int(K[0, 0] * point3D_new[0] / point3D_new[2] + K[0, 2])
                v = int(K[1, 1] * point3D_new[1] / point3D_new[2] + K[1, 2])
                
                if 0 <= u < w and 0 <= v < h:
                    new_img[v, u] = img[y, x]
        return new_img
        
    def wrap_img(self, img, depth_map, K, R, T, target_point=None):
        h, w = img.shape[:2]
        # Generate grid of coordinates (x, y)
        x, y = np.meshgrid(np.arange(w), np.arange(h))
        ones = np.ones_like(x)

        # Flatten and stack to get homogeneous coordinates
        homogeneous_coordinates = np.stack((x.flatten(), y.flatten(), ones.flatten()), axis=1).T

        # Inverse intrinsic matrix
        K_inv = np.linalg.inv(K)

        # Inverse rotation and translation
        R_inv = R.T
        T_inv = -R_inv @ T

        # Project to 3D using depth map
        world_coordinates = K_inv @ homogeneous_coordinates
        world_coordinates *= depth_map.flatten()

        # Apply inverse transformation
        transformed_world_coordinates = R_inv @ world_coordinates + T_inv.reshape(-1, 1)

        # Project back to 2D
        valid = transformed_world_coordinates[2, :] > 0
        projected_2D = K @ transformed_world_coordinates
        projected_2D /= projected_2D[2, :]

        # Initialize map_x and map_y
        map_x = np.zeros((h, w), dtype=np.float32)
        map_y = np.zeros((h, w), dtype=np.float32)

        # Assign valid projection values to map_x and map_y
        map_x.flat[valid] = projected_2D[0, valid]
        map_y.flat[valid] = projected_2D[1, valid]

        # Perform the warping
        wrapped_img = cv2.remap(img, map_x, map_y, interpolation=cv2.INTER_LINEAR)
        if target_point is None:
            return wrapped_img
        else:
            target_point = (map_x[int(target_point[1]), int(target_point[0])], map_y[int(target_point[1]), int(target_point[0])])
            target_point = tuple(max(0, min(511, x)) for x in target_point)
            return wrapped_img, target_point

    def get_low_high_frequent_tensors(self, x, threshold=4):
        dtype = x.dtype
        x = x.type(torch.float32)

        # FFT
        x_freq = fft.fftn(x, dim=(-2, -1))
        x_freq = fft.fftshift(x_freq, dim=(-2, -1))
        B,C,H,W = x_freq.shape
        mask = torch.ones((B, C, H, W)).to(x.device)

        crow, ccol = H // 2, W //2
        mask[..., crow - threshold:crow + threshold, ccol - threshold:ccol + threshold] = 0 # low 0 high 1
        x_freq_high = x_freq * mask
        x_freq_low = x_freq * (1 - mask)

        x_freq_high = fft.ifftshift(x_freq_high, dim=(-2, -1))
        x_high = fft.ifftn(x_freq_high, dim=(-2, -1)).real
        x_high = x_high.type(dtype)

        x_freq_low = fft.ifftshift(x_freq_low, dim=(-2, -1))
        x_low = fft.ifftn(x_freq_low, dim=(-2, -1)).real
        x_low = x_low.type(dtype)
        return x_high, x_low, x_freq_high, x_freq_low, mask
    
    def combine_low_and_high(self, freq_low, freq_high, mask):
        freq = freq_high * mask + freq_low * (1-mask)
        freq = fft.ifftshift(freq, dim=(-2, -1))
        x = fft.ifftn(freq, dim=(-2, -1)).real
        return x


    def wrap_img_tensor_w_fft(self, img_tensor, depth_tensor, 
                        theta_z=0, theta_x=0, theta_y=-10, T=[0,0,-2], threshold=4):
        _, img_tensor, high_freq, low_freq, fft_mask = self.get_low_high_frequent_tensors(img_tensor, threshold)

        intrinsic_matrix = np.array([[1000, 0, img_tensor.shape[-1]/2],
                    [0, 1000, img_tensor.shape[-2]/2],
                    [0, 0, 1]])  # Example intrinsic matrix
        ori_size = None
        if depth_tensor.shape[-1] != img_tensor.shape[-1]:
            scale = depth_tensor.shape[-1] / img_tensor.shape[-1]
            ori_size = (img_tensor.shape[-2], img_tensor.shape[-1])
            img_tensor_ori = img_tensor.clone()
            # img_tensor = F.interpolate(img_tensor, size=(depth_tensor.shape[-2], depth_tensor.shape[-1]))
            depth_tensor = F.interpolate(depth_tensor.unsqueeze(0).unsqueeze(0), size=ori_size, mode='bilinear').squeeze().to(torch.float16)
            intrinsic_matrix[0,0] /= scale
            intrinsic_matrix[1,1] /= scale
        rotation_matrix = self.calc_R(theta_z=theta_z, theta_x=theta_x, theta_y=theta_y)
        translation_vector = np.array(T)  # Translation vector to shift camera to the right

        h,w = img_tensor.shape[2:]

        xy_src = np.mgrid[0:h, 0:w].reshape(2, -1)

        xy_src_homogeneous = np.vstack((xy_src, np.ones((1, xy_src.shape[1]))))

        # Convert to torch tensors
        xy_src_homogeneous_tensor = torch.tensor(xy_src_homogeneous, dtype=torch.float16, device=img_tensor.device)

        # Compute the coordinates in the world frame
        xy_world = torch.inverse(torch.tensor(intrinsic_matrix, dtype=torch.float32, device=img_tensor.device)).to(torch.float16) @ xy_src_homogeneous_tensor
        xy_world = xy_world * depth_tensor.view(1, -1)

        # Compute the coordinates in the new camera frame
        xy_new_cam = torch.inverse(torch.tensor(rotation_matrix, dtype=torch.float32, device=img_tensor.device)).to(torch.float16) @ (xy_world - torch.tensor(translation_vector, dtype=torch.float16, device=img_tensor.device).view(3,1))

        # Compute the coordinates in the new image
        xy_dst_homogeneous = torch.tensor(intrinsic_matrix, dtype=torch.float16, device=img_tensor.device) @ xy_new_cam
        xy_dst = xy_dst_homogeneous[:2, :] / xy_dst_homogeneous[2, :]

        # Reshape to a 2D grid and normalize to [-1, 1]
        xy_dst = xy_dst.reshape(2, h, w)
        xy_dst = (xy_dst - torch.tensor([[w/2.0], [h/2.0]], dtype=torch.float16, device=img_tensor.device).unsqueeze(-1)) / torch.tensor([[w/2.0], [h/2.0]], dtype=torch.float16, device=img_tensor.device).unsqueeze(-1)
        xy_dst = torch.flip(xy_dst, [0])
        xy_dst = xy_dst.permute(1, 2, 0)

        # Perform the warping
        wrapped_img = F.grid_sample(img_tensor, xy_dst.to(torch.float16)[None], align_corners=True, mode='bilinear', padding_mode='reflection')
        wrapped_freq = fft.fftn(wrapped_img, dim=(-2, -1))
        wrapped_freq = fft.fftshift(wrapped_freq, dim=(-2, -1))
        wrapped_img = self.combine_low_and_high(wrapped_freq, high_freq, fft_mask)
        return wrapped_img
    
    def wrap_img_tensor_w_fft_ext(self, img_tensor, depth_tensor, K,R,T, threshold=4):
        _, img_tensor, high_freq, low_freq, fft_mask = self.get_low_high_frequent_tensors(img_tensor, threshold)

        ori_size = None

        if depth_tensor.shape[-1] != img_tensor.shape[-1]:
            scale = depth_tensor.shape[-1] / img_tensor.shape[-1]
            ori_size = (img_tensor.shape[-2], img_tensor.shape[-1])
            # img_tensor = F.interpolate(img_tensor, size=(depth_tensor.shape[-2], depth_tensor.shape[-1]))
            depth_tensor = F.interpolate(depth_tensor.unsqueeze(0).unsqueeze(0), size=ori_size, mode='bilinear').squeeze().to(torch.float16)
            intrinsic = copy.deepcopy(K)
            intrinsic = K / scale
            intrinsic[2,2] = 1

        h,w = img_tensor.shape[2:]

        xy_src = np.mgrid[0:h, 0:w].reshape(2, -1)

        xy_src_homogeneous = np.vstack((xy_src, np.ones((1, xy_src.shape[1]))))

        # Convert to torch tensors
        xy_src_homogeneous_tensor = torch.tensor(xy_src_homogeneous, dtype=img_tensor.dtype, device=img_tensor.device)

        # Compute the coordinates in the world frame
        # xy_world = torch.inverse(torch.tensor(K, dtype=torch.float32, device=img_tensor.device)).to(torch.float16) @ xy_src_homogeneous_tensor
        xy_world = torch.tensor(np.linalg.inv(intrinsic)).to(img_tensor.dtype).to(img_tensor.device) @ xy_src_homogeneous_tensor
        xy_world = xy_world * depth_tensor.view(1, -1)

        # Compute the coordinates in the new camera frame
        xy_new_cam = torch.inverse(torch.tensor(R, dtype=torch.float32, device=img_tensor.device)).to(img_tensor.dtype) @ (xy_world - torch.tensor(T, dtype=img_tensor.dtype, device=img_tensor.device).view(3,1))

        # Compute the coordinates in the new image
        xy_dst_homogeneous = torch.tensor(intrinsic, dtype=img_tensor.dtype, device=img_tensor.device) @ xy_new_cam
        xy_dst = xy_dst_homogeneous[:2, :] / xy_dst_homogeneous[2, :]

        # Reshape to a 2D grid and normalize to [-1, 1]
        xy_dst = xy_dst.reshape(2, h, w)
        xy_dst = (xy_dst - torch.tensor([[w/2.0], [h/2.0]], dtype=img_tensor.dtype, device=img_tensor.device).unsqueeze(-1)) / torch.tensor([[w/2.0], [h/2.0]], dtype=img_tensor.dtype, device=img_tensor.device).unsqueeze(-1)
        xy_dst = torch.flip(xy_dst, [0])
        xy_dst = xy_dst.permute(1, 2, 0)

        # Perform the warping
        wrapped_img = F.grid_sample(img_tensor, xy_dst.to(img_tensor.dtype)[None], align_corners=True, mode='bilinear', padding_mode='reflection')
        wrapped_freq = fft.fftn(wrapped_img, dim=(-2, -1))
        wrapped_freq = fft.fftshift(wrapped_freq, dim=(-2, -1))
        wrapped_img = self.combine_low_and_high(wrapped_freq, high_freq, fft_mask)
        return wrapped_img

    def wrap_img_tensor_w_fft_matrix(self, img_tensor, depth_tensor, 
                        theta_z=0, theta_x=0, theta_y=-10, T=[0,0,-2], threshold=4):
        _, img_tensor, high_freq, low_freq, fft_mask = self.get_low_high_frequent_tensors(img_tensor, threshold)

        intrinsic_matrix = np.array([[1000, 0, img_tensor.shape[-1]/2],
                    [0, 1000, img_tensor.shape[-2]/2],
                    [0, 0, 1]])  # Example intrinsic matrix
        ori_size = None
        if depth_tensor.shape[-1] != img_tensor.shape[-1]:
            scale = depth_tensor.shape[-1] / img_tensor.shape[-1]
            ori_size = (img_tensor.shape[-2], img_tensor.shape[-1])
            img_tensor_ori = img_tensor.clone()
            # img_tensor = F.interpolate(img_tensor, size=(depth_tensor.shape[-2], depth_tensor.shape[-1]))
            depth_tensor = F.interpolate(depth_tensor.unsqueeze(0).unsqueeze(0), size=ori_size, mode='bilinear').squeeze().to(torch.float16)
            intrinsic_matrix[0,0] /= scale
            intrinsic_matrix[1,1] /= scale
        rotation_matrix = self.calc_R(theta_z=theta_z, theta_x=theta_x, theta_y=theta_y)
        translation_vector = np.array(T)  # Translation vector to shift camera to the right

        h,w = img_tensor.shape[2:]

        xy_src = np.mgrid[0:h, 0:w].reshape(2, -1)

        xy_src_homogeneous = np.vstack((xy_src, np.ones((1, xy_src.shape[1]))))

        # Convert to torch tensors
        xy_src_homogeneous_tensor = torch.tensor(xy_src_homogeneous, dtype=torch.float16, device=img_tensor.device)

        # Compute the coordinates in the world frame
        xy_world = torch.inverse(torch.tensor(intrinsic_matrix, dtype=torch.float32, device=img_tensor.device)).to(torch.float16) @ xy_src_homogeneous_tensor
        xy_world = xy_world * depth_tensor.view(1, -1)

        # Compute the coordinates in the new camera frame
        xy_new_cam = torch.inverse(torch.tensor(rotation_matrix, dtype=torch.float32, device=img_tensor.device)).to(torch.float16) @ (xy_world - torch.tensor(translation_vector, dtype=torch.float16, device=img_tensor.device).view(3,1))

        # Compute the coordinates in the new image
        xy_dst_homogeneous = torch.tensor(intrinsic_matrix, dtype=torch.float16, device=img_tensor.device) @ xy_new_cam
        xy_dst = xy_dst_homogeneous[:2, :] / xy_dst_homogeneous[2, :]

        # Reshape to a 2D grid and normalize to [-1, 1]
        xy_dst = xy_dst.reshape(2, h, w)
        xy_dst = (xy_dst - torch.tensor([[w/2.0], [h/2.0]], dtype=torch.float16, device=img_tensor.device).unsqueeze(-1)) / torch.tensor([[w/2.0], [h/2.0]], dtype=torch.float16, device=img_tensor.device).unsqueeze(-1)
        xy_dst = torch.flip(xy_dst, [0])
        xy_dst = xy_dst.permute(1, 2, 0)

        # Perform the warping
        wrapped_img = F.grid_sample(img_tensor, xy_dst.to(torch.float16)[None], align_corners=True, mode='bilinear', padding_mode='reflection')
        wrapped_freq = fft.fftn(wrapped_img, dim=(-2, -1))
        wrapped_freq = fft.fftshift(wrapped_freq, dim=(-2, -1))
        wrapped_img = self.combine_low_and_high(wrapped_freq, high_freq, fft_mask)


        return wrapped_img
    
    
    def wrap_img_tensor(self, img_tensor, depth_tensor, 
                        theta_z=0, theta_x=0, theta_y=-10, T=[0,0,-2]):
        intrinsic_matrix = np.array([[1000, 0, img_tensor.shape[-1]/2],
                    [0, 1000, img_tensor.shape[-2]/2],
                    [0, 0, 1]])  # Example intrinsic matrix
        ori_size = None
        if depth_tensor.shape[-1] != img_tensor.shape[-1]:
            scale = depth_tensor.shape[-1] / img_tensor.shape[-1]
            ori_size = (img_tensor.shape[-2], img_tensor.shape[-1])
            img_tensor_ori = img_tensor.clone()
            # img_tensor = F.interpolate(img_tensor, size=(depth_tensor.shape[-2], depth_tensor.shape[-1]))
            depth_tensor = F.interpolate(depth_tensor.unsqueeze(0).unsqueeze(0), size=ori_size, mode='bilinear').squeeze().to(torch.float16)
            intrinsic_matrix[0,0] /= scale
            intrinsic_matrix[1,1] /= scale
        rotation_matrix = self.calc_R(theta_z=theta_z, theta_x=theta_x, theta_y=theta_y)
        translation_vector = np.array(T)  # Translation vector to shift camera to the right

        h,w = img_tensor.shape[2:]

        xy_src = np.mgrid[0:h, 0:w].reshape(2, -1)

        xy_src_homogeneous = np.vstack((xy_src, np.ones((1, xy_src.shape[1]))))

        # Convert to torch tensors
        xy_src_homogeneous_tensor = torch.tensor(xy_src_homogeneous, dtype=torch.float16, device=img_tensor.device)

        # Compute the coordinates in the world frame
        xy_world = torch.inverse(torch.tensor(intrinsic_matrix, dtype=torch.float32, device=img_tensor.device)).to(torch.float16) @ xy_src_homogeneous_tensor
        xy_world = xy_world * depth_tensor.view(1, -1)

        # Compute the coordinates in the new camera frame
        xy_new_cam = torch.inverse(torch.tensor(rotation_matrix, dtype=torch.float32, device=img_tensor.device)).to(torch.float16) @ (xy_world - torch.tensor(translation_vector, dtype=torch.float16, device=img_tensor.device).view(3,1))

        # Compute the coordinates in the new image
        xy_dst_homogeneous = torch.tensor(intrinsic_matrix, dtype=torch.float16, device=img_tensor.device) @ xy_new_cam
        xy_dst = xy_dst_homogeneous[:2, :] / xy_dst_homogeneous[2, :]

        # Reshape to a 2D grid and normalize to [-1, 1]
        xy_dst = xy_dst.reshape(2, h, w)
        xy_dst = (xy_dst - torch.tensor([[w/2.0], [h/2.0]], dtype=torch.float16, device=img_tensor.device).unsqueeze(-1)) / torch.tensor([[w/2.0], [h/2.0]], dtype=torch.float16, device=img_tensor.device).unsqueeze(-1)
        xy_dst = torch.flip(xy_dst, [0])
        xy_dst = xy_dst.permute(1, 2, 0)

        # Perform the warping
        wrapped_img = F.grid_sample(img_tensor, xy_dst.to(torch.float16)[None], align_corners=True, mode='bilinear')



        return wrapped_img

    @torch.no_grad()
    def __call__(self, img_array, theta_z=0, theta_x=0, theta_y=-10, T=[0,0,-2]):
        img_depth = self.transform({"image": img_array})["image"]

        # compute
        prediction = self.process(
            self.device,
            self.model,
            self.model_type,
            img_depth,
            (self.net_w, self.net_h),
            img_array.shape[1::-1],
            optimize=False,
            use_camera=False,
        )

        depth = self.prediction2depth(prediction)

        # img = img_array.copy()
        # img = img / 2. + 0.5
        K = np.array([[1000, 0, img_array.shape[1]/2],
                    [0, 1000, img_array.shape[0]/2],
                    [0, 0, 1]])  # Example intrinsic matrix

        R = self.calc_R(theta_z=theta_z, theta_x=theta_x, theta_y=theta_y)
        T = np.array(T)  # Translation vector to shift camera to the right

        # new_img = self.render_new_view(img_array, depth, R, T, K)
        new_img = self.wrap_img(img_array, depth, K, R, T)

        mask = np.all(new_img == [0,0,0], axis=2).astype(np.uint8) * 255
        mask = 255 - mask
        return new_img, mask, depth