model/model_diff_data.py

"""
model to deal with shapenet inputs and other datasets such as Behave and ProciGen
the model takes a different data dictionary in forward function
"""
import inspect
from typing import Optional
import numpy as np

import torch
import torch.nn.functional as F
from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
from diffusers.schedulers.scheduling_ddim import DDIMScheduler
from diffusers.schedulers.scheduling_pndm import PNDMScheduler
from pytorch3d.implicitron.dataset.data_loader_map_provider import FrameData
from pytorch3d.renderer.cameras import CamerasBase
from pytorch3d.structures import Pointclouds
from torch import Tensor
from tqdm import tqdm
from pytorch3d.renderer import PerspectiveCameras
from pytorch3d.datasets.r2n2.utils import BlenderCamera


from .model import ConditionalPointCloudDiffusionModel
from .model_utils import get_num_points


class ConditionalPCDiffusionShapenet(ConditionalPointCloudDiffusionModel):
    def forward(self, batch, mode: str = 'train', **kwargs):
        """
        take a batch of data from ShapeNet
        """
        images = torch.stack(batch['images'], 0).to('cuda')
        masks = torch.stack(batch['masks'], 0).to('cuda')
        pc = Pointclouds([x.to('cuda') for x in batch['pclouds']])
        camera = BlenderCamera(
            torch.stack(batch['R']),
            torch.stack(batch['T']),
            torch.stack(batch['K']), device='cuda'
        )

        if mode == 'train':
            return self.forward_train(
                pc=pc,
                camera=camera,
                image_rgb=images,
                mask=masks,

                **kwargs)
        elif mode == 'sample':
            num_points = kwargs.pop('num_points', get_num_points(pc))
            return self.forward_sample(
                num_points=num_points,
                camera=camera,
                image_rgb=images,
                mask=masks,
                gt_pc=pc,
                **kwargs)
        else:
            raise NotImplementedError()


class ConditionalPCDiffusionBehave(ConditionalPointCloudDiffusionModel):
    "diffusion model for Behave dataset"
    def forward(self, batch, mode: str = 'train', **kwargs):
        images = torch.stack(batch['images'], 0).to('cuda')
        masks = torch.stack(batch['masks'], 0).to('cuda')
        pc = self.get_input_pc(batch)
        camera = PerspectiveCameras(
            R=torch.stack(batch['R']),
            T=torch.stack(batch['T']),
            K=torch.stack(batch['K']),
            device='cuda',
            in_ndc=True
        )
        grid_df = torch.stack(batch['grid_df'], 0).to('cuda') if 'grid_df' in batch else None
        num_points = kwargs.pop('num_points', get_num_points(pc))
        if mode == 'train':
            return self.forward_train(
                pc=pc,
                camera=camera,
                image_rgb=images,
                mask=masks,
                grid_df=grid_df,
                **kwargs)
        elif mode == 'sample':
            return self.forward_sample(
                num_points=num_points,
                camera=camera,
                image_rgb=images,
                mask=masks,
                gt_pc=pc,
                **kwargs)
        else:
            raise NotImplementedError()

    def get_input_pc(self, batch):
        pc = Pointclouds([x.to('cuda') for x in batch['pclouds']])
        return pc


class ConditionalPCDiffusionSeparateSegm(ConditionalPCDiffusionBehave):
    "a separate model to predict binary labels, the final segmentation model"
    def __init__(self,
                 beta_start: float,
                 beta_end: float,
                 beta_schedule: str,
                 point_cloud_model: str,
                 point_cloud_model_embed_dim: int,
                 **kwargs,  # projection arguments
                 ):
        super(ConditionalPCDiffusionSeparateSegm, self).__init__(beta_start, beta_end, beta_schedule,
                                                                 point_cloud_model,
                                                                 point_cloud_model_embed_dim, **kwargs)
        # add a separate model to predict binary label
        from .point_cloud_transformer_model import PointCloudTransformerModel, PointCloudModel

        self.binary_model = PointCloudTransformerModel(
            num_layers=1, # XH: use the default color model number of layers
            model_type=point_cloud_model, # pvcnn
            embed_dim=point_cloud_model_embed_dim, # save as pc shape model
            in_channels=self.in_channels,
            out_channels=1,
        )
        self.binary_training_noise_std = kwargs.get("binary_training_noise_std", 0.1)

        # re-initialize point cloud model
        assert self.predict_binary
        self.point_cloud_model = PointCloudModel(
            model_type=point_cloud_model,
            embed_dim=point_cloud_model_embed_dim,
            in_channels=self.in_channels,
            out_channels=self.out_channels - 1,  # not predicting binary from this anymore
            voxel_resolution_multiplier=kwargs.get('voxel_resolution_multiplier', 1)
        )

    def forward_train(
        self,
        pc: Pointclouds,
        camera: Optional[CamerasBase],
        image_rgb: Optional[Tensor],
        mask: Optional[Tensor],
        return_intermediate_steps: bool = False,
        **kwargs
    ):
        # first run shape forward, then binary label forward
        assert not return_intermediate_steps
        assert self.predict_binary
        loss_shape = super(ConditionalPCDiffusionSeparateSegm, self).forward_train(pc,
                                                                                   camera,
                                                                                   image_rgb,
                                                                                   mask,
                                                                                   return_intermediate_steps,
                                                                                   **kwargs)

        # binary label forward
        x_0 = self.point_cloud_to_tensor(pc, normalize=True, scale=True)
        x_points, x_colors = x_0[:, :, :3], x_0[:, :, 3:]

        # Add noise to points.
        x_input = x_points + torch.randn_like(x_points) * self.binary_training_noise_std # std=0.1
        x_input = self.get_input_with_conditioning(x_input, camera=camera,
                                                   image_rgb=image_rgb, mask=mask, t=None)

        # Forward
        pred_segm = self.binary_model(x_input)

        # use compressed bits
        df_grid = kwargs.get('grid_df', None).unsqueeze(1)  # (B, 1, resz, resy, resx)
        points = x_points.clone().detach() / self.scale_factor * 2  # , normalize to [-1, 1]
        points[:, :, 0], points[:, :, 2] = points[:, :, 2].clone(), points[:, :,0].clone()  # swap, make sure clone is used!
        points = points.unsqueeze(1).unsqueeze(1)  # (B,1, 1, N, 3)
        with torch.no_grad():
            df_interp = F.grid_sample(df_grid, points, padding_mode='border', align_corners=True).squeeze(1).squeeze(1)  # (B, 1, 1, 1, N)
        binary_label = df_interp[:, 0] > 0.5  # (B, 1, N)

        binary_pred = torch.sigmoid(pred_segm.squeeze(-1))  # add a sigmoid layer
        loss_binary = F.mse_loss(binary_pred, binary_label.float().squeeze(1).squeeze(1)) * self.lw_binary
        loss = loss_shape + loss_binary

        return loss, torch.tensor([loss_shape, loss_binary])

    def reverse_step(self, extra_step_kwargs, scheduler, t, x_t, x_t_input, **kwargs):
        "return (B, N, 4), the 4-th channel is binary label"
        B = x_t.shape[0]
        # Forward
        noise_pred = self.point_cloud_model(x_t_input, t.reshape(1).expand(B))
        if self.consistent_center:
            assert self.dm_pred_type != 'sample', 'incompatible dm predition type!'
            # suggested by the CCD-3DR paper
            noise_pred = noise_pred - torch.mean(noise_pred, dim=1, keepdim=True)
        # Step: make sure only update the shape (first 3 channels)
        x_t = scheduler.step(noise_pred, t, x_t[:, :, :3], **extra_step_kwargs).prev_sample
        if self.consistent_center:
            x_t = x_t - torch.mean(x_t, dim=1, keepdim=True)

        # also add binary prediction
        if kwargs.get('inference_binary', False):
            pred_segm = self.binary_model(x_t_input)
        else:
            pred_segm = torch.zeros_like(x_t[:, :, 0:1])

        x_t = torch.cat([x_t, torch.sigmoid(pred_segm)], -1)

        return x_t

    def get_coord_feature(self, x_t):
        x_t_input = [x_t[:, :, :3]]
        return x_t_input

    def tensor_to_point_cloud(self, x: Tensor, /, denormalize: bool = False, unscale: bool = False):
        """
        take binary label into account
        :param self:
        :param x: (B, N, 4), the 4th channel is the binary segmentation, 1-human, 0-object
        :param denormalize: denormalize the per-point colors, from pc2
        :param unscale: undo point scaling, from pc2
        :return: pc with point colors if predict binary label or per-point color
        """
        points = x[:, :, :3] / (self.scale_factor if unscale else 1)
        if self.predict_color:
            colors = self.denormalize(x[:, :, 3:]) if denormalize else x[:, :, 3:]
            return Pointclouds(points=points, features=colors)
        else:
            if self.predict_binary:
                assert x.shape[2] == 4
                # add color to predicted binary labels
                is_hum = x[:, :, 3] > 0.5
                features = []
                for mask in is_hum:
                    color = torch.zeros_like(x[0, :, :3]) + torch.tensor([0.5, 1.0, 0]).to(x.device)
                    color[mask, :] = torch.tensor([0.05, 1.0, 1.0]).to(x.device) # human is light blue, object light green
                    features.append(color)
            else:
                assert x.shape[2] == 3
                features = None
            return Pointclouds(points=points, features=features)