# author: Justus Schock (justus.schock@rwth-aachen.de)
import numpy as np
import torch
import os
from torch.utils.cpp_extension import load as load_cpp
[docs]class ShapeLayer(torch.nn.Module):
"""
Wrapper to compine Python and C++ Implementation under Single API
"""
def __init__(self, shapes, use_cpp=False):
"""
Parameters
----------
shapes : np.ndarray
the actual shape components
use_cpp : bool
whether or not to use the (experimental) C++ Implementation
"""
super().__init__()
if use_cpp:
self._layer = _ShapeLayerCpp(shapes)
else:
self._layer = _ShapeLayerPy(shapes)
[docs] def forward(self, shape_params: torch.Tensor):
"""
Forwards parameters to Python or C++ Implementation
Parameters
----------
shape_params : :class:`torch.Tensor`
parameters for shape ensembling
Returns
-------
:class:`torch.Tensor`
Ensempled Shape
"""
return self._layer(shape_params)
@property
def num_params(self):
"""
Property to access these layer's parameters
Returns
-------
int
number of parameters
"""
return self._layer.num_params
[docs]class _ShapeLayerPy(torch.nn.Module):
"""
Python Implementation of Shape Layer
"""
def __init__(self, shapes):
"""
Parameters
----------
shapes : np.ndarray
eigen shapes (obtained by PCA)
"""
super().__init__()
self.register_buffer("_shape_mean", torch.from_numpy(
shapes[0]).float().unsqueeze(0))
components = []
for i, _shape in enumerate(shapes[1:]):
components.append(torch.from_numpy(
_shape).float().unsqueeze(0))
component_tensor = torch.cat(components).unsqueeze(0)
self.register_buffer("_shape_components", component_tensor)
[docs] def forward(self, shape_params: torch.Tensor):
"""
Ensemble shape from parameters
Parameters
----------
shape_params : :class:`torch.Tensor`
shape parameters
Returns
-------
:class:`torch.Tensor`
ensembled shape
"""
shapes = getattr(self, "_shape_mean").clone()
shapes = shapes.expand(shape_params.size(0), *shapes.size()[1:])
components = getattr(self, "_shape_components")
components = components.expand(shape_params.size(0),
*components.size()[1:])
weighted_components = components.mul(
shape_params.expand_as(components))
shapes = shapes.add(weighted_components.sum(dim=1))
return shapes
@property
def num_params(self):
"""
Property to access these layer's parameters
Returns
-------
int
number of parameters
"""
return getattr(self, "_shape_components").size(1)
[docs]class _ShapeLayerCpp(torch.nn.Module):
"""
C++ Implementation of Shape Layer
"""
def __init__(self, shapes, verbose=True):
"""
Parameters
----------
shapes : np.ndarray
eigen shapes (obtained by PCA)
"""
super().__init__()
self.register_buffer("_shape_mean",
torch.from_numpy(shapes[0]).float().unsqueeze(0))
components = []
for i, _shape in enumerate(shapes[1:]):
components.append(torch.from_numpy(_shape).float().unsqueeze(0))
component_tensor = torch.cat(components).unsqueeze(0)
self.register_buffer("_shape_components", component_tensor)
self._func = load_cpp("shape_function",
sources=[os.path.join(os.path.split(__file__)[0],
"shape_layer.cpp")],
verbose=verbose)
[docs] def forward(self, shape_params: torch.Tensor):
"""
Ensemble shape from parameters
Parameters
----------
shape_params : :class:`torch.Tensor`
shape parameters
Returns
-------
:class:`torch.Tensor`
ensembled shape
"""
shapes = self._func.forward(shape_params,
getattr(self, "_shape_mean"),
getattr(self, "_shape_components"))
return shapes
@property
def num_params(self):
"""
Property to access these layer's parameters
Returns
-------
int
number of parameters
"""
return getattr(self, "_shape_components").size(1)