File size: 8,307 Bytes
162943d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
from __future__ import division, print_function
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.autograd import Function

import voxelize_cuda


class VoxelizationFunction(Function):
    """
    Definition of differentiable voxelization function
    Currently implemented only for cuda Tensors
    """
    @staticmethod
    def forward(ctx, smpl_vertices, smpl_face_center, smpl_face_normal,
                smpl_vertex_code, smpl_face_code, smpl_tetrahedrons,
                volume_res, sigma, smooth_kernel_size):
        """
        forward pass
        Output format: (batch_size, z_dims, y_dims, x_dims, channel_num) 
        """
        assert (smpl_vertices.size()[1] == smpl_vertex_code.size()[1])
        assert (smpl_face_center.size()[1] == smpl_face_normal.size()[1])
        assert (smpl_face_center.size()[1] == smpl_face_code.size()[1])
        ctx.batch_size = smpl_vertices.size()[0]
        ctx.volume_res = volume_res
        ctx.sigma = sigma
        ctx.smooth_kernel_size = smooth_kernel_size
        ctx.smpl_vertex_num = smpl_vertices.size()[1]
        ctx.device = smpl_vertices.device

        smpl_vertices = smpl_vertices.contiguous()
        smpl_face_center = smpl_face_center.contiguous()
        smpl_face_normal = smpl_face_normal.contiguous()
        smpl_vertex_code = smpl_vertex_code.contiguous()
        smpl_face_code = smpl_face_code.contiguous()
        smpl_tetrahedrons = smpl_tetrahedrons.contiguous()

        occ_volume = torch.cuda.FloatTensor(ctx.batch_size, ctx.volume_res,
                                            ctx.volume_res,
                                            ctx.volume_res).fill_(0.0)
        semantic_volume = torch.cuda.FloatTensor(ctx.batch_size,
                                                 ctx.volume_res,
                                                 ctx.volume_res,
                                                 ctx.volume_res, 3).fill_(0.0)
        weight_sum_volume = torch.cuda.FloatTensor(ctx.batch_size,
                                                   ctx.volume_res,
                                                   ctx.volume_res,
                                                   ctx.volume_res).fill_(1e-3)

        # occ_volume [B, volume_res, volume_res, volume_res]
        # semantic_volume [B, volume_res, volume_res, volume_res, 3]
        # weight_sum_volume [B, volume_res, volume_res, volume_res]

        occ_volume, semantic_volume, weight_sum_volume = voxelize_cuda.forward_semantic_voxelization(
            smpl_vertices, smpl_vertex_code, smpl_tetrahedrons, occ_volume,
            semantic_volume, weight_sum_volume, sigma)

        return semantic_volume


class Voxelization(nn.Module):
    """
    Wrapper around the autograd function VoxelizationFunction
    """

    def __init__(self, smpl_vertex_code, smpl_face_code, smpl_face_indices,
                 smpl_tetraderon_indices, volume_res, sigma,
                 smooth_kernel_size, batch_size, device):
        super(Voxelization, self).__init__()
        assert (len(smpl_face_indices.shape) == 2)
        assert (len(smpl_tetraderon_indices.shape) == 2)
        assert (smpl_face_indices.shape[1] == 3)
        assert (smpl_tetraderon_indices.shape[1] == 4)

        self.volume_res = volume_res
        self.sigma = sigma
        self.smooth_kernel_size = smooth_kernel_size
        self.batch_size = batch_size
        self.device = device

        self.smpl_vertex_code = smpl_vertex_code
        self.smpl_face_code = smpl_face_code
        self.smpl_face_indices = smpl_face_indices
        self.smpl_tetraderon_indices = smpl_tetraderon_indices

    def update_param(self, batch_size, smpl_tetra):

        self.batch_size = batch_size
        self.smpl_tetraderon_indices = smpl_tetra

        smpl_vertex_code_batch = np.tile(self.smpl_vertex_code,
                                         (self.batch_size, 1, 1))
        smpl_face_code_batch = np.tile(self.smpl_face_code,
                                       (self.batch_size, 1, 1))
        smpl_face_indices_batch = np.tile(self.smpl_face_indices,
                                          (self.batch_size, 1, 1))
        smpl_tetraderon_indices_batch = np.tile(self.smpl_tetraderon_indices,
                                                (self.batch_size, 1, 1))

        smpl_vertex_code_batch = torch.from_numpy(
            smpl_vertex_code_batch).contiguous().to(self.device)
        smpl_face_code_batch = torch.from_numpy(
            smpl_face_code_batch).contiguous().to(self.device)
        smpl_face_indices_batch = torch.from_numpy(
            smpl_face_indices_batch).contiguous().to(self.device)
        smpl_tetraderon_indices_batch = torch.from_numpy(
            smpl_tetraderon_indices_batch).contiguous().to(self.device)

        self.register_buffer('smpl_vertex_code_batch', smpl_vertex_code_batch)
        self.register_buffer('smpl_face_code_batch', smpl_face_code_batch)
        self.register_buffer('smpl_face_indices_batch',
                             smpl_face_indices_batch)
        self.register_buffer('smpl_tetraderon_indices_batch',
                             smpl_tetraderon_indices_batch)

    def forward(self, smpl_vertices):
        """
        Generate semantic volumes from SMPL vertices
        """
        assert (smpl_vertices.size()[0] == self.batch_size)
        self.check_input(smpl_vertices)
        smpl_faces = self.vertices_to_faces(smpl_vertices)
        smpl_tetrahedrons = self.vertices_to_tetrahedrons(smpl_vertices)
        smpl_face_center = self.calc_face_centers(smpl_faces)
        smpl_face_normal = self.calc_face_normals(smpl_faces)
        smpl_surface_vertex_num = self.smpl_vertex_code_batch.size()[1]
        smpl_vertices_surface = smpl_vertices[:, :smpl_surface_vertex_num, :]
        vol = VoxelizationFunction.apply(smpl_vertices_surface,
                                         smpl_face_center, smpl_face_normal,
                                         self.smpl_vertex_code_batch,
                                         self.smpl_face_code_batch,
                                         smpl_tetrahedrons, self.volume_res,
                                         self.sigma, self.smooth_kernel_size)
        return vol.permute((0, 4, 1, 2, 3))  # (bzyxc --> bcdhw)

    def vertices_to_faces(self, vertices):
        assert (vertices.ndimension() == 3)
        bs, nv = vertices.shape[:2]
        device = vertices.device
        face = self.smpl_face_indices_batch + (
            torch.arange(bs, dtype=torch.int32).to(device) * nv)[:, None, None]
        vertices_ = vertices.reshape((bs * nv, 3))
        return vertices_[face.long()]

    def vertices_to_tetrahedrons(self, vertices):
        assert (vertices.ndimension() == 3)
        bs, nv = vertices.shape[:2]
        device = vertices.device
        tets = self.smpl_tetraderon_indices_batch + (
            torch.arange(bs, dtype=torch.int32).to(device) * nv)[:, None, None]
        vertices_ = vertices.reshape((bs * nv, 3))
        return vertices_[tets.long()]

    def calc_face_centers(self, face_verts):
        assert len(face_verts.shape) == 4
        assert face_verts.shape[2] == 3
        assert face_verts.shape[3] == 3
        bs, nf = face_verts.shape[:2]
        face_centers = (face_verts[:, :, 0, :] + face_verts[:, :, 1, :] +
                        face_verts[:, :, 2, :]) / 3.0
        face_centers = face_centers.reshape((bs, nf, 3))
        return face_centers

    def calc_face_normals(self, face_verts):
        assert len(face_verts.shape) == 4
        assert face_verts.shape[2] == 3
        assert face_verts.shape[3] == 3
        bs, nf = face_verts.shape[:2]
        face_verts = face_verts.reshape((bs * nf, 3, 3))
        v10 = face_verts[:, 0] - face_verts[:, 1]
        v12 = face_verts[:, 2] - face_verts[:, 1]
        normals = F.normalize(torch.cross(v10, v12), eps=1e-5)
        normals = normals.reshape((bs, nf, 3))
        return normals

    def check_input(self, x):
        if x.device == 'cpu':
            raise TypeError('Voxelization module supports only cuda tensors')
        if x.type() != 'torch.cuda.FloatTensor':
            raise TypeError(
                'Voxelization module supports only float32 tensors')