ProposalNetwork/utils/spaces.py

import numpy as np
import torch
from cubercnn import util

'''
coordinate system is assumed to have origin in the upper left
(0,0) _________________(N,0)
|  
|    
| 
|
|
(0,M)
'''
"""
class Cube:
    '''
    3D box in the format [c1, c2, c3, w, h, l, R]

    Args:
        c1: The x coordinate of the center of the box.
        c2: The y coordinate of the center of the box.
        c3: The z coordinate of the center of the box.
        w: The width of the box in meters.
        h: The height of the box in meters.
        l: The length of the box in meters.
        R: The 3D rotation matrix of the box.
    ```

                      _____________________ 
                    /|                    /|
                   / |                   / |
                  /  |                  /  |
                 /___|_________________/   |
                |    |                 |   | h
                |    |                 |   |
                |    |                 |   |
                |    |   (c1,c2,c3)    |   |
                |    |_________________|___|
                |   /                  |   /
                |  /                   |  /
                | /                    | / l
                |/_____________________|/
                            w             
    ```
    '''
    def __init__(self,tensor: torch.Tensor, R: torch.Tensor, score=None, label=None) -> None:
        self.tensor = tensor
        self.center = tensor[:3]
        self.dimensions = tensor[3:6]
        self.rotation = R

        # score and label are meant as auxiliary information
        self.score = score
        self.label = label

    def get_cube(self):
        color = [c/255.0 for c in util.get_color()]
        return util.mesh_cuboid(torch.cat((self.center,self.dimensions)), self.rotation, color=color)
    
    def get_all_corners(self):
        '''wrap ``util.get_cuboid_verts_faces``
        
        Returns:
            verts: the 3D vertices of the cuboid in camera space'''
        verts, _ = util.get_cuboid_verts_faces(torch.cat((self.center,self.dimensions)), self.rotation)
        return verts
    
    def get_bube_corners(self,K) -> torch.Tensor:
        cube_corners = self.get_all_corners()
        cube_corners = torch.mm(K, cube_corners.t()).t()
        return cube_corners[:,:2]/cube_corners[:,2].unsqueeze(1)
    
    def get_volume(self) -> float:
        return self.dimensions.prod().item()

    
    def __repr__(self) -> str:
        return f'Cube({self.center}, {self.dimensions}, {self.rotation})'
    
    def to_device(self, device):
        '''
        Move all tensors of the instantiated class to the specified device.

        Args:
            device: The device to move the tensors to (e.g., 'cuda', 'cpu').
        '''
        self.tensor = self.tensor.to(device)
        self.center = self.center.to(device)
        self.dimensions = self.dimensions.to(device)
        self.rotation = self.rotation.to(device)
        return self
"""

class Cubes:
    '''
    3D boxes in the format [[c1, c2, c3, w, h, l, R1...R9]]

    inspired by `detectron2.structures.Boxes`

    Args:
        tensor: torch.tensor(
            c1: The x coordinates of the center of the boxes.
            c2: The y coordinates of the center of the boxes.
            c3: The z coordinates of the center of the boxes.
            w: The width of the boxes in meters.
            h: The height of the boxes in meters.
            l: The length of the boxes in meters.
            R: The flattened 3D rotation matrix of the boxes (i.e. the rows are next to each other).
            )
            of shape (N, 15).
    ```
                      _____________________ 
                    /|                    /|
                   / |                   / |
                  /  |                  /  |
                 /___|_________________/   |
                |    |                 |   | h
                |    |                 |   |
                |    |                 |   |
                |    |   (c1,c2,c3)    |   |
                |    |_________________|___|
                |   /                  |   /
                |  /                   |  /
                | /                    | / l
                |/_____________________|/
                            w             
    ```
    '''
    def __init__(self,tensor: torch.Tensor, scores=None, labels=None) -> None:

        # score and label are meant as auxiliary information
        if scores is not None:
            assert scores.ndim == 2, f"scores.shape must be (n_instances, n_proposals), but was {scores.shape}" 
        self.scores = scores
        self.labels = labels

        if not isinstance(tensor, torch.Tensor):
            if not isinstance(tensor, np.ndarray):
                tensor = np.asarray(tensor)
            tensor = torch.as_tensor(tensor, dtype=torch.float32, device=torch.device("cpu"))
        else:
            tensor = tensor.to(torch.float32)
        if tensor.numel() == 0:
            tensor = tensor.reshape((-1, 15)).to(dtype=torch.float32)
        self.tensor = tensor
        if self.tensor.dim() == 1:
            self.tensor = self.tensor.unsqueeze(0)
        if self.tensor.dim() == 2:
            self.tensor = self.tensor.unsqueeze(0)

    @property
    def centers(self):
        return self.tensor[:, :, :3]
    
    @property
    def dimensions(self):
        return self.tensor[:, :, 3:6]
    
    @property
    def rotations(self):
        shape = self.tensor.shape
        return self.tensor[:, :, 6:].reshape(shape[0],shape[1], 3, 3)
    
    @property
    def device(self):
        return self.tensor.device
    
    @property
    def num_instances(self):
        return self.tensor.shape[0]
    
    @property
    def shape(self):
        return self.tensor.shape

    def clone(self) -> "Cubes":
        """
        Clone the Cubes.

        Returns:
            Cubes
        """
        return Cubes(self.tensor.clone())
    

    def get_cubes(self):
        color = [c/255.0 for c in util.get_color()]
        return util.mesh_cuboid(torch.cat((self.centers.squeeze(0),self.dimensions.squeeze(0)),dim=1), self.rotations.squeeze(0), color=color)
        
    
    def get_all_corners(self):
        '''wrap ``util.get_cuboid_verts_faces``
        
        Returns:
            verts: the 3D vertices of the cuboid in camera space'''

        verts_list = []
        for i in range(self.num_instances):
            verts_next_instance, _ = util.get_cuboid_verts_faces(self.tensor[i, :, :6], self.rotations[i])
            verts_list.append(verts_next_instance)
        verts = torch.stack(verts_list, dim=0)

        return verts
    
    def get_cuboids_verts_faces(self):
        '''wrap ``util.get_cuboid_verts_faces``
        
        Returns:
            verts: the 3D vertices of the cuboid in camera space
            faces: the faces of the cuboid in camera space'''

        verts_list = []
        faces_list = []
        for i in range(self.num_instances):
            verts_next_instance, faces = util.get_cuboid_verts_faces(self.tensor[i, :, :6], self.rotations[i])
            verts_list.append(verts_next_instance)
            faces_list.append(faces)
        verts = torch.stack(verts_list, dim=0)
        faces = torch.stack(faces_list, dim=0)

        return verts, faces
    
    def get_bube_corners(self, K, clamp:tuple=None) -> torch.Tensor:
        '''This assumes that all the cubes have the same camera intrinsic matrix K

        clamp is a typically the image shape (width, height) to truncate the boxes to image frame, this avoids huge projected boxes
        Returns:
            num_instances x N x 8 x 2'''
        cube_corners = self.get_all_corners() # num_instances x N x 8 x 3
        num_prop = cube_corners.shape[1]
        cube_corners = cube_corners.reshape(self.num_instances * num_prop, 8, 3)
        K_repeated = K.repeat(self.num_instances * num_prop,1,1)
        cube_corners = torch.matmul(K_repeated, cube_corners.transpose(2,1))
        cube_corners = cube_corners[:, :2, :]/cube_corners[:, 2, :].unsqueeze(-2)
        cube_corners = cube_corners.transpose(2,1)
        cube_corners = cube_corners.reshape(self.num_instances, num_prop, 8, 2)

        # we must clamp and then stack, otherwise the gradient is fucked
        if clamp is not None:
            x = torch.clamp(cube_corners[..., 0], int(-clamp[0]/2+1), int(clamp[0]-1+clamp[0]))
            y = torch.clamp(cube_corners[..., 1], int(-clamp[1]/2+1), int(clamp[1]-1+clamp[1]))
        cube_corners = torch.stack((x, y), dim=-1)
        
        return cube_corners # num_instances x num_proposals x 8 x 2
    
    def get_volumes(self) -> float:
        return self.get_dimensions().prod(1).item()
    
    def __len__(self) -> int:
        return self.tensor.shape[0]

    def __repr__(self) -> str:
        return f'Cubes({self.tensor})'
    
    def to(self, device: torch.device):
        # Cubes are assumed float32 and does not support to(dtype)
        if isinstance(self.scores, torch.Tensor):
            self.scores = self.scores.to(device=device)
        if isinstance(self.labels, torch.Tensor):
            self.labels = self.labels.to(device=device)
        return Cubes(self.tensor.to(device=device), self.scores, self.labels)
    
    def __getitem__(self, item) -> "Cubes":
        """
        Args:
            item: int, slice, or a BoolTensor

        Returns:
            Cubes: Create a new :class:`Cubes` by indexing.

        The following usage are allowed:

        1. `new_cubes = cubes[3]`: return a `Cubes` which contains only one box.
        2. `new_cubes = cubes[2:10]`: return a slice of cubes.
        3. `new_cubes = cubes[vector]`, where vector is a torch.BoolTensor
           with `length = len(cubes)`. Nonzero elements in the vector will be selected.

        Note that the returned Cubes might share storage with this Cubes,
        subject to Pytorch's indexing semantics.
        """
        if isinstance(item, int):
            prev_n_prop = self.tensor.shape[1]
            return Cubes(self.tensor[item].view(1, prev_n_prop, -1))
        elif isinstance(item, tuple):
            return Cubes(self.tensor[item[0],item[1]].view(1, 1, -1))
        b = self.tensor[item]
        assert b.dim() == 2, "Indexing on Cubes with {} failed to return a matrix!".format(item)
        return Cubes(b)
    

    @classmethod
    def cat(cls, cubes_list: list["Cubes"]) -> "Cubes":
        """
        Concatenates a list of Cubes into a single Cubes

        Arguments:
            cubes_list (list[Cubes])

        Returns:
            Cubes: the concatenated Cubes
        """
        assert isinstance(cubes_list, (list, tuple))
        if len(cubes_list) == 0:
            return cls(torch.empty(0))
        assert all([isinstance(box, Cubes) for box in cubes_list])

        # use torch.cat (v.s. layers.cat) so the returned cubes never share storage with input
        cat_cubes = cls(torch.cat([b.tensor for b in cubes_list], dim=0))
        return cat_cubes
    
    @torch.jit.unused
    def __iter__(self):
        """
        Yield a cube as a Tensor of shape (15,) at a time.
        """
        yield from self.tensor

    def split(self, split_size: int, dim=1) -> tuple["Cubes"]:
        """same behaviour as torch.split, return a tuple of chunksize Cubes"""
        return tuple(Cubes(x) for x in self.tensor.split(split_size, dim=dim))
    
    def reshape(self, *args) -> "Cubes":
        """
        Returns:
            Cubes: reshaped Cubes
        """
        return Cubes(self.tensor.reshape(*args), self.scores, self.labels)