models/dsine.py

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from models.submodules import Encoder, ConvGRU, UpSampleBN, UpSampleGN, RayReLU, \
                                convex_upsampling, get_unfold, get_prediction_head, \
                                INPUT_CHANNELS_DICT
from utils.rotation import axis_angle_to_matrix


class Decoder(nn.Module):
    def __init__(self, output_dims, B=5, NF=2048, BN=False, downsample_ratio=8):
        super(Decoder, self).__init__()
        input_channels = INPUT_CHANNELS_DICT[B]
        output_dim, feature_dim, hidden_dim = output_dims
        features = bottleneck_features = NF
        self.downsample_ratio = downsample_ratio

        UpSample = UpSampleBN if BN else UpSampleGN 
        self.conv2 = nn.Conv2d(bottleneck_features + 2, features, kernel_size=1, stride=1, padding=0)
        self.up1 = UpSample(skip_input=features // 1 + input_channels[1] + 2, output_features=features // 2, align_corners=False)
        self.up2 = UpSample(skip_input=features // 2 + input_channels[2] + 2, output_features=features // 4, align_corners=False)

        # prediction heads
        i_dim = features // 4
        h_dim = 128
        self.normal_head = get_prediction_head(i_dim+2, h_dim, output_dim)
        self.feature_head = get_prediction_head(i_dim+2, h_dim, feature_dim)
        self.hidden_head = get_prediction_head(i_dim+2, h_dim, hidden_dim)

    def forward(self, features, uvs):
        _, _, x_block2, x_block3, x_block4 = features[4], features[5], features[6], features[8], features[11]
        uv_32, uv_16, uv_8 = uvs

        x_d0 = self.conv2(torch.cat([x_block4, uv_32], dim=1))
        x_d1 = self.up1(x_d0, torch.cat([x_block3, uv_16], dim=1))
        x_feat = self.up2(x_d1, torch.cat([x_block2, uv_8], dim=1))
        x_feat = torch.cat([x_feat, uv_8], dim=1)

        normal = self.normal_head(x_feat)
        normal = F.normalize(normal, dim=1)
        f = self.feature_head(x_feat)
        h = self.hidden_head(x_feat)
        return normal, f, h


class DSINE(nn.Module):
    def __init__(self):
        super(DSINE, self).__init__()
        self.downsample_ratio = 8
        self.ps = 5           # patch size
        self.num_iter = 5     # num iterations

        # define encoder
        self.encoder = Encoder(B=5, pretrained=True)

        # define decoder
        self.output_dim = output_dim = 3
        self.feature_dim = feature_dim = 64
        self.hidden_dim = hidden_dim = 64
        self.decoder = Decoder([output_dim, feature_dim, hidden_dim], B=5, NF=2048, BN=False)

        # ray direction-based ReLU
        self.ray_relu = RayReLU(eps=1e-2)

        # pixel_coords (1, 3, H, W)
        # NOTE: this is set to some arbitrarily high number, 
        # if your input is 2000+ pixels wide/tall, increase these values
        h = 2000
        w = 2000
        pixel_coords = np.ones((3, h, w)).astype(np.float32)
        x_range = np.concatenate([np.arange(w).reshape(1, w)] * h, axis=0)
        y_range = np.concatenate([np.arange(h).reshape(h, 1)] * w, axis=1)
        pixel_coords[0, :, :] = x_range + 0.5
        pixel_coords[1, :, :] = y_range + 0.5
        self.pixel_coords = torch.from_numpy(pixel_coords).unsqueeze(0)

        # define ConvGRU cell
        self.gru = ConvGRU(hidden_dim=hidden_dim, input_dim=feature_dim+2, ks=self.ps)

        # padding used during NRN
        self.pad = (self.ps - 1) // 2

        # prediction heads
        self.prob_head = get_prediction_head(self.hidden_dim+2, 64, self.ps*self.ps)   # weights assigned for each nghbr pixel 
        self.xy_head = get_prediction_head(self.hidden_dim+2, 64, self.ps*self.ps*2)   # rotation axis for each nghbr pixel
        self.angle_head = get_prediction_head(self.hidden_dim+2, 64, self.ps*self.ps)  # rotation angle for each nghbr pixel

        # prediction heads - weights used for upsampling the coarse resolution output
        self.up_prob_head = get_prediction_head(self.hidden_dim+2, 64, 9 * self.downsample_ratio * self.downsample_ratio)

    def get_ray(self, intrins, H, W, orig_H, orig_W, return_uv=False):
        B, _, _ = intrins.shape
        fu = intrins[:, 0, 0][:,None,None] * (W / orig_W)
        cu = intrins[:, 0, 2][:,None,None] * (W / orig_W)
        fv = intrins[:, 1, 1][:,None,None] * (H / orig_H)
        cv = intrins[:, 1, 2][:,None,None] * (H / orig_H)

        # (B, 2, H, W)
        ray = self.pixel_coords[:, :, :H, :W].repeat(B, 1, 1, 1)
        ray[:, 0, :, :] = (ray[:, 0, :, :] - cu) / fu
        ray[:, 1, :, :] = (ray[:, 1, :, :] - cv) / fv

        if return_uv:
            return ray[:, :2, :, :]
        else:
            return F.normalize(ray, dim=1)

    def upsample(self, h, pred_norm, uv_8):
        up_mask = self.up_prob_head(torch.cat([h, uv_8], dim=1))
        up_pred_norm = convex_upsampling(pred_norm, up_mask, self.downsample_ratio)
        up_pred_norm = F.normalize(up_pred_norm, dim=1)
        return up_pred_norm

    def refine(self, h, feat_map, pred_norm, intrins, orig_H, orig_W, uv_8, ray_8):
        B, C, H, W = pred_norm.shape
        fu = intrins[:, 0, 0][:,None,None,None] * (W / orig_W)  # (B, 1, 1, 1)
        cu = intrins[:, 0, 2][:,None,None,None] * (W / orig_W)
        fv = intrins[:, 1, 1][:,None,None,None] * (H / orig_H)
        cv = intrins[:, 1, 2][:,None,None,None] * (H / orig_H)

        h_new = self.gru(h, feat_map)

        # get nghbr prob (B, 1, ps*ps, h, w)
        nghbr_prob = self.prob_head(torch.cat([h_new, uv_8], dim=1)).unsqueeze(1)
        nghbr_prob = torch.sigmoid(nghbr_prob)

        # get nghbr normals (B, 3, ps*ps, h, w)
        nghbr_normals = get_unfold(pred_norm, ps=self.ps, pad=self.pad)

        # get nghbr xy (B, 2, ps*ps, h, w)
        nghbr_xys = self.xy_head(torch.cat([h_new, uv_8], dim=1))
        nghbr_xs, nghbr_ys = torch.split(nghbr_xys, [self.ps*self.ps, self.ps*self.ps], dim=1)
        nghbr_xys = torch.cat([nghbr_xs.unsqueeze(1), nghbr_ys.unsqueeze(1)], dim=1)        
        nghbr_xys = F.normalize(nghbr_xys, dim=1)

        # get nghbr theta (B, 1, ps*ps, h, w)
        nghbr_angle = self.angle_head(torch.cat([h_new, uv_8], dim=1)).unsqueeze(1)
        nghbr_angle = torch.sigmoid(nghbr_angle) * np.pi

        # get nghbr pixel coord (1, 3, ps*ps, h, w)
        nghbr_pixel_coord = get_unfold(self.pixel_coords[:, :, :H, :W], ps=self.ps, pad=self.pad)

        # nghbr axes (B, 3, ps*ps, h, w)
        nghbr_axes = torch.zeros_like(nghbr_normals)

        du_over_fu = nghbr_xys[:, 0, ...] / fu                                      # (B, ps*ps, h, w)
        dv_over_fv = nghbr_xys[:, 1, ...] / fv                                      # (B, ps*ps, h, w)

        term_u = (nghbr_pixel_coord[:, 0, ...] + nghbr_xys[:, 0, ...] - cu) / fu    # (B, ps*ps, h, w)
        term_v = (nghbr_pixel_coord[:, 1, ...] + nghbr_xys[:, 1, ...] - cv) / fv    # (B, ps*ps, h, w)

        nx = nghbr_normals[:, 0, ...]                                               # (B, ps*ps, h, w)
        ny = nghbr_normals[:, 1, ...]                                               # (B, ps*ps, h, w)
        nz = nghbr_normals[:, 2, ...]                                               # (B, ps*ps, h, w)

        nghbr_delta_z_num = - (du_over_fu * nx + dv_over_fv * ny)
        nghbr_delta_z_denom = (term_u * nx + term_v * ny + nz)
        nghbr_delta_z_denom[torch.abs(nghbr_delta_z_denom) < 1e-8] = 1e-8 * torch.sign(nghbr_delta_z_denom[torch.abs(nghbr_delta_z_denom) < 1e-8])
        nghbr_delta_z = nghbr_delta_z_num / nghbr_delta_z_denom

        nghbr_axes[:, 0, ...] = du_over_fu + nghbr_delta_z * term_u
        nghbr_axes[:, 1, ...] = dv_over_fv + nghbr_delta_z * term_v
        nghbr_axes[:, 2, ...] = nghbr_delta_z
        nghbr_axes = F.normalize(nghbr_axes, dim=1)                                 # (B, 3, ps*ps, h, w)

        # make sure axes are all valid
        invalid = torch.sum(torch.logical_or(torch.isnan(nghbr_axes), torch.isinf(nghbr_axes)).float(), dim=1) > 0.5    # (B, ps*ps, h, w)
        nghbr_axes[:, 0, ...][invalid] = 0.0
        nghbr_axes[:, 1, ...][invalid] = 0.0
        nghbr_axes[:, 2, ...][invalid] = 0.0

        # nghbr_axes_angle (B, 3, ps*ps, h, w)
        nghbr_axes_angle = nghbr_axes * nghbr_angle
        nghbr_axes_angle = nghbr_axes_angle.permute(0, 2, 3, 4, 1)  # (B, ps*ps, h, w, 3)
        nghbr_R = axis_angle_to_matrix(nghbr_axes_angle)            # (B, ps*ps, h, w, 3, 3)

        # (B, 3, ps*ps, h, w)
        nghbr_normals_rot = torch.bmm(
            nghbr_R.reshape(B * self.ps * self.ps * H * W, 3, 3),
            nghbr_normals.permute(0, 2, 3, 4, 1).reshape(B * self.ps * self.ps * H * W, 3).unsqueeze(-1)
        ).reshape(B, self.ps*self.ps, H, W, 3, 1).squeeze(-1).permute(0, 4, 1, 2, 3)        # (B, 3, ps*ps, h, w)
        nghbr_normals_rot = F.normalize(nghbr_normals_rot, dim=1)

        # ray ReLU
        nghbr_normals_rot = torch.cat([
            self.ray_relu(nghbr_normals_rot[:, :, i, :, :], ray_8).unsqueeze(2) 
            for i in range(nghbr_normals_rot.size(2))
            ], dim=2)

        # (B, 1, ps*ps, h, w) * (B, 3, ps*ps, h, w)
        pred_norm = torch.sum(nghbr_prob * nghbr_normals_rot, dim=2)    # (B, C, H, W)
        pred_norm = F.normalize(pred_norm, dim=1)
    
        up_mask = self.up_prob_head(torch.cat([h_new, uv_8], dim=1))
        up_pred_norm = convex_upsampling(pred_norm, up_mask, self.downsample_ratio)
        up_pred_norm = F.normalize(up_pred_norm, dim=1)

        return h_new, pred_norm, up_pred_norm

    
    def forward(self, img, intrins=None):
        # Step 1. encoder
        features = self.encoder(img)

        # Step 2. get uv encoding
        B, _, orig_H, orig_W = img.shape
        intrins[:, 0, 2] += 0.5
        intrins[:, 1, 2] += 0.5
        uv_32 = self.get_ray(intrins, orig_H//32, orig_W//32, orig_H, orig_W, return_uv=True)
        uv_16 = self.get_ray(intrins, orig_H//16, orig_W//16, orig_H, orig_W, return_uv=True)
        uv_8 = self.get_ray(intrins, orig_H//8, orig_W//8, orig_H, orig_W, return_uv=True)
        ray_8 = self.get_ray(intrins, orig_H//8, orig_W//8, orig_H, orig_W)

        # Step 3. decoder - initial prediction
        pred_norm, feat_map, h = self.decoder(features, uvs=(uv_32, uv_16, uv_8))
        pred_norm = self.ray_relu(pred_norm, ray_8)

        # Step 4. add ray direction encoding
        feat_map = torch.cat([feat_map, uv_8], dim=1)

        # iterative refinement
        up_pred_norm = self.upsample(h, pred_norm, uv_8)
        pred_list = [up_pred_norm]
        for i in range(self.num_iter):
            h, pred_norm, up_pred_norm = self.refine(h, feat_map, 
                                                     pred_norm.detach(), 
                                                     intrins, orig_H, orig_W, uv_8, ray_8)
            pred_list.append(up_pred_norm)
        return pred_list