Skip to content

3D RGB Residuals #13

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
erincar wants to merge 16 commits into
base: master
Choose a base branch
from
Open

3D RGB Residuals #13

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
880e587
Add RGBLoss3D template
erincar Feb 9, 2021
990765c
Mark potential usage of RGBLoss3D
erincar Feb 9, 2021
69ef9c8
Add new notebook
erincar Feb 9, 2021
4e9efa0
Add RGB residual method, Use in combined ICP
erincar Feb 16, 2021
100f296
Use pointcloud colors in RGB optimization
erincar Feb 16, 2021
5254f06
Prioritize dense RGB optimier in combined ICP
erincar Feb 16, 2021
2b0a19a
Use different nearest neighbor modes for dense & rgb
erincar Feb 16, 2021
125036a
Implement support for color coefficient optimization
erincar Feb 16, 2021
dc8ad50
Add more color parameter support
erincar Feb 16, 2021
c33d5b3
Add RGB residuals to combined loss function
erincar Feb 16, 2021
8f053b8
Test combined ICP with RGB, color weight = 1000
erincar Feb 23, 2021
152ecee
Test combined ICP with RGB, color weight = 10000
erincar Feb 23, 2021
fc138fb
Test combined ICP with RGB, sparse = 1000, dense = 10000
erincar Feb 23, 2021
e2e8c58
Add color residual weighting, w=100
erincar Feb 23, 2021
2e10ead
Crank color residual weight up to 200
erincar Feb 23, 2021
52082c0
Merge branch 'master' into erinc/rgb-residuals
erincar Mar 1, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
111 changes: 90 additions & 21 deletions face_reconstruction/optim/bfm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,8 +35,8 @@

- Parameters (BFMOptimizationParameters):
Syntactic sugar. Provides an interface to translate between a list of parameters (theta) and an object of
meaningful attributes, e.g., divided into shape coefficients, expression coefficients and camera pose

meaningful attributes, e.g., divided into shape coefficients, expression coefficients, color coefficients
and camera pose
"""


Expand All @@ -54,9 +54,11 @@ def __init__(self,
bfm: BaselFaceModel,
n_params_shape,
n_params_expression,
n_params_color,
fix_camera_pose=False,
weight_shape_params=1.0,
weight_expression_params=1.0,
weight_color_params=1.0,
rotation_mode='lie'):
"""

Expand All @@ -70,20 +72,27 @@ def __init__(self,
Specifies that only the first `n_params_expression` parameters of the expression model will be optimized for.
These are the parameters that have the biggest impact on the face model.
The remaining coefficients will be held constant to 0.
:param n_params_color:
Specifies that only the first `n_params_color` parameters of the expression model will be optimized for.
These are the parameters that have the biggest impact on the face model.
The remaining coefficients will be held constant to 0.
:param fix_camera_pose:
Whether the camera pose should be optimized for
:param weight_shape_params:
Specifies how much more changing a shape coefficient parameter will impact the loss
:param weight_expression_params:
Specifies how much more changing an expression coefficient parameter will impact the loss
:param weight_color_params:
Specifies how much more changing a color coefficient parameter will impact the loss
"""
self.bfm = bfm
self.n_params_shape = n_params_shape
self.n_params_expression = n_params_expression
self.n_params_color = 0 # Currently, optimizing for color is not supported
self.n_params_color = n_params_color
self.fix_camera_pose = fix_camera_pose
self.weight_shape_params = weight_shape_params
self.weight_expression_params = weight_expression_params
self.weight_color_params = weight_color_params

assert rotation_mode in ['quaternion', 'lie'], f'Rotation mode has to be either lie or quaternion. ' \
f'You gave {rotation_mode}'
Expand All @@ -96,12 +105,14 @@ def __init__(self,
lower_bounds = []
lower_bounds.extend([-float('inf') for _ in range(n_params_shape)])
lower_bounds.extend([-float('inf') for _ in range(n_params_expression)])
lower_bounds.extend([-float('inf') for _ in range(n_params_color)])
lower_bounds.extend([-1, -1, -1, -1, -float('inf'), -float('inf'), -float('inf')])
self.lower_bounds = np.array(lower_bounds)

upper_bounds = []
upper_bounds.extend([float('inf') for _ in range(n_params_shape)])
upper_bounds.extend([float('inf') for _ in range(n_params_expression)])
upper_bounds.extend([float('inf') for _ in range(n_params_color)])
upper_bounds.extend([1, 1, 1, 1, float('inf'), float('inf'), float('inf')])
self.upper_bounds = np.array(upper_bounds)

Expand All @@ -126,12 +137,14 @@ def create_optimization_context(loss, initial_params, max_nfev=100, verbose=2, x
def create_parameters(self,
shape_coefficients: np.ndarray = None,
expression_coefficients: np.ndarray = None,
color_coefficients: np.ndarray = None,
camera_pose: np.ndarray = None
):
return BFMOptimizationParameters(
self,
shape_coefficients=shape_coefficients,
expression_coefficients=expression_coefficients,
color_coefficients=expression_coefficients,
camera_pose=camera_pose)

def create_parameters_from_other(self, parameters):
Expand All @@ -144,6 +157,7 @@ def create_sparse_loss(self,
fixed_camera_pose: np.ndarray = None,
fixed_shape_coefficients: np.ndarray = None,
fixed_expression_coefficients: np.ndarray = None,
fixed_color_coefficients: np.ndarray = None,
regularization_strength: float = None
):
return SparseOptimizationLoss(
Expand All @@ -154,6 +168,7 @@ def create_sparse_loss(self,
fixed_camera_pose=fixed_camera_pose,
fixed_shape_coefficients=fixed_shape_coefficients,
fixed_expression_coefficients=fixed_expression_coefficients,
fixed_color_coefficients=fixed_color_coefficients,
regularization_strength=regularization_strength)

def create_sparse_loss_3d(self,
Expand Down Expand Up @@ -187,15 +202,21 @@ def create_combined_loss_3d(self,
nearest_neighbor_mode: NearestNeighborMode,
distance_type: DistanceType,
weight_sparse_term: float = 1,
weight_color_term: float = 1,
regularization_strength: float = None,
pointcloud_normals: np.ndarray = None
pointcloud_normals: np.ndarray = None,
pointcloud_colors: np.ndarray = None,
):
return CombinedLoss3D(self, bfm_landmark_indices=bfm_landmark_indices,
img_landmark_points_3d=img_landmark_points_3d,
pointcloud=pointcloud, nearest_neighbors=nearest_neighbors,
nearest_neighbor_mode=nearest_neighbor_mode,
distance_type=distance_type, weight_sparse_term=weight_sparse_term,
regularization_strength=regularization_strength, pointcloud_normals=pointcloud_normals)
distance_type=distance_type,
weight_sparse_term=weight_sparse_term,
weight_color_term=weight_color_term,
regularization_strength=regularization_strength,
pointcloud_normals=pointcloud_normals,
pointcloud_colors=pointcloud_colors)

def create_sparse_keyframe_loss(self, bfm_landmark_indices_list: List[np.ndarray],
img_landmark_points_3d_list: List[np.ndarray],
Expand Down Expand Up @@ -230,7 +251,7 @@ class BFMOptimizationContext:
Encapsulates the context of a single optimization run. Needed in order to hold the initial parameters as
the theta list that is communicated to the optimizer only contains the parameters that are to be optimized.
Fixed parameters, i.e., those specified as initial parameters that are outside the n_shape_coefficients or
n_expression_coefficients bounds, can then be obtained from this context wrapper.
n_expression_coefficients or n_color_coefficients bounds, can then be obtained from this context wrapper.
"""

def __init__(self, loss, initial_params, max_nfev=100, verbose=2, x_scale='jac'):
Expand Down Expand Up @@ -318,12 +339,14 @@ def _apply_params_to_model(self, theta):

shape_coefficients = parameters.shape_coefficients
expression_coefficients = parameters.expression_coefficients
color_coefficients = parameters.color_coefficients
camera_pose = parameters.camera_pose

face_mesh = self.optimization_manager.bfm.draw_sample(
shape_coefficients=shape_coefficients,
expression_coefficients=expression_coefficients,
color_coefficients=[0 for _ in range(self.optimization_manager.n_color_coefficients)])
color_coefficients=color_coefficients,
)
bfm_vertices = add_column(np.array(face_mesh.vertices), 1)
bfm_vertices = camera_pose @ bfm_vertices.T
return bfm_vertices.T, face_mesh
Expand Down Expand Up @@ -372,6 +395,7 @@ def __init__(
fixed_camera_pose: np.ndarray = None,
fixed_shape_coefficients: np.ndarray = None,
fixed_expression_coefficients: np.ndarray = None,
fixed_color_coefficients: np.ndarray = None,
regularization_strength=None):
"""
:param optimization_manager:
Expand All @@ -394,6 +418,7 @@ def __init__(

self.fixed_shape_coefficients = fixed_shape_coefficients
self.fixed_expression_coefficients = fixed_expression_coefficients
self.fixed_color_coefficients = fixed_color_coefficients

def loss(self, theta, *args, **kwargs):
parameters = self.create_parameters_from_theta(theta)
Expand All @@ -408,6 +433,11 @@ def loss(self, theta, *args, **kwargs):
else:
expression_coefficients = self.fixed_expression_coefficients

if self.fixed_color_coefficients is None:
color_coefficients = parameters.color_coefficients
else:
color_coefficients = self.fixed_color_coefficients

if self.optimization_manager.fix_camera_pose:
camera_pose = self.fixed_camera_pose
else:
Expand All @@ -416,7 +446,8 @@ def loss(self, theta, *args, **kwargs):
face_mesh = self.optimization_manager.bfm.draw_sample(
shape_coefficients=shape_coefficients,
expression_coefficients=expression_coefficients,
color_coefficients=[0 for _ in range(self.optimization_manager.n_color_coefficients)])
color_coefficients=color_coefficients,
)
landmark_points = np.array(face_mesh.vertices)[self.bfm_landmark_indices]
face_landmark_pixels = self.renderer.project_points(camera_pose, landmark_points)
residuals = face_landmark_pixels - self.img_landmark_pixels
Expand Down Expand Up @@ -538,8 +569,10 @@ def __init__(self, optimization_manager: BFMOptimization,
nearest_neighbor_mode: NearestNeighborMode,
distance_type: DistanceType,
weight_sparse_term: float = 1,
weight_color_term: float = 1,
regularization_strength: float = None,
pointcloud_normals: np.ndarray = None):
pointcloud_normals: np.ndarray = None,
pointcloud_colors: np.ndarray = None):
super(CombinedLoss3D, self).__init__(optimization_manager, regularization_strength)
self.bfm_landmark_indices = bfm_landmark_indices
self.img_landmark_points_3d = img_landmark_points_3d
Expand All @@ -548,7 +581,9 @@ def __init__(self, optimization_manager: BFMOptimization,
self.nearest_neighbor_mode = nearest_neighbor_mode
self.distance_type = distance_type
self.pointcloud_normals = pointcloud_normals
self.pointcloud_colors = pointcloud_colors
self.weight_sparse_term = weight_sparse_term
self.weight_color_term = weight_color_term

def loss(self, theta, *args, **kwargs):
bfm_vertices, face_mesh = self._apply_params_to_model(theta)
Expand Down Expand Up @@ -589,6 +624,16 @@ def loss(self, theta, *args, **kwargs):
landmark_points = bfm_vertices[self.bfm_landmark_indices]
residuals.extend(self.weight_sparse_term * (landmark_points[:, :3] - self.img_landmark_points_3d))

# Color residuals
if self.pointcloud_colors is not None:
if self.nearest_neighbor_mode == NearestNeighborMode.FACE_VERTICES:
mesh_colors = np.array(face_mesh.colors)
pointcloud_colors = self.pointcloud_colors[self.nearest_neighbors]
elif self.nearest_neighbor_mode == NearestNeighborMode.POINTCLOUD:
mesh_colors = np.array(face_mesh.colors)[self.nearest_neighbors]
pointcloud_colors = self.pointcloud_colors
residuals.extend(self.weight_color_term * (mesh_colors - pointcloud_colors))

residuals = np.array(residuals).reshape(-1)
if self.regularization_strength is not None:
regularization_terms = self._compute_regularization_terms(
Expand Down Expand Up @@ -724,6 +769,7 @@ def __init__(self,
optimization_manager: BFMOptimization,
shape_coefficients: np.ndarray,
expression_coefficients: np.ndarray,
color_coefficients: np.ndarray,
camera_pose: np.ndarray):
"""
Defines all the parameters that will be optimized for
Expand All @@ -733,6 +779,8 @@ def __init__(self,
The part of the parameters that describes the shape coefficients
:param expression_coefficients:
The part of the parameters that describes the expression coefficients
:param color_coefficients:
The part of the parameters that describes the color coefficients
:param camera_pose:
The part of the parameters that describes the 4x4 camera pose matrix
"""
Expand All @@ -743,6 +791,7 @@ def __init__(self,
n_color_coefficients = optimization_manager.n_color_coefficients
n_params_shape = optimization_manager.n_params_shape
n_params_expression = optimization_manager.n_params_expression
n_params_color = optimization_manager.n_params_color

assert shape_coefficients is not None or n_params_shape == 0, "If n_params_shape > 0 then shape coefficients have to be provided"
if shape_coefficients is None:
Expand All @@ -751,14 +800,21 @@ def __init__(self,
assert expression_coefficients is not None or n_params_expression == 0, "If n_params_expression > 0 then expression coefficients have to be provided"
if expression_coefficients is None:
expression_coefficients = []
# Shape and expression coefficients are multiplied by their weight to enforce that changing them

assert color_coefficients is not None or n_params_color == 0, "If n_params_color > 0 then color coefficients have to be provided"
if color_coefficients is None:
color_coefficients = []
# Shape, expression & color coefficients are multiplied by their weight to enforce that changing them
# will have a higher impact depending on the weight
self.shape_coefficients = np.hstack(
[shape_coefficients,
np.zeros((n_shape_coefficients - len(shape_coefficients)))])
self.expression_coefficients = np.hstack([expression_coefficients,
np.zeros((n_expression_coefficients - len(expression_coefficients)))])
self.color_coefficients = np.zeros(n_color_coefficients)
self.color_coefficients = np.hstack([
color_coefficients,
np.zeros((n_color_coefficients - len(color_coefficients))),
])

assert camera_pose is not None or optimization_manager.fix_camera_pose, "Camera pose may only be None if it is fixed"
if camera_pose is not None:
Expand All @@ -773,6 +829,7 @@ def from_theta(optimization_manager: Union[BFMOptimization, BFMOptimizationConte
Contains a list of parameters that are interpreted as follows.
The 1st `n_shape_params` are shape coefficients
The next `n_expression_params` are expression coefficients
The next `n_color_params` are expression coefficients
The final 7 parameters are the quaternion defining the camera rotation (4 params) and the translation (3 params)
"""
context = None
Expand All @@ -782,22 +839,30 @@ def from_theta(optimization_manager: Union[BFMOptimization, BFMOptimizationConte

n_params_shape = optimization_manager.n_params_shape
n_params_expression = optimization_manager.n_params_expression
n_params_color = optimization_manager.n_params_color

if context is None:
# No access to initial or fixed parameters
# Just reconstruct coefficients from what is there in the theta list
shape_coefficients = theta[:n_params_shape] * optimization_manager.weight_shape_params
expression_coefficients = theta[n_params_shape:n_params_shape + n_params_expression] \
* optimization_manager.weight_expression_params
start, end = 0, n_params_shape
shape_coefficients = theta[start:end] * optimization_manager.weight_shape_params
start, end = end, end + n_params_expression
expression_coefficients = theta[start:end] * optimization_manager.weight_expression_params
start, end = end, end + n_params_color
color_coefficients = theta[start:end] * optimization_manager.weight_color_params
else:
# Access to initial or fixed parameters
# Combine initial parameters and what is there in the theta list
start, end = 0, n_params_shape
shape_coefficients = context.initial_params.shape_coefficients
shape_coefficients[:n_params_shape] = theta[:n_params_shape] * optimization_manager.weight_shape_params
shape_coefficients[:n_params_shape] = theta[start:end] * optimization_manager.weight_shape_params
start, end = end, end + n_params_expression
expression_coefficients = context.initial_params.expression_coefficients
expression_coefficients[:n_params_expression] = theta[n_params_shape:n_params_shape + n_params_expression] \
* optimization_manager.weight_expression_params
i = n_params_shape + n_params_expression
expression_coefficients[:n_params_expression] = theta[start:end] * optimization_manager.weight_expression_params
start, end = end, end + n_params_color
color_coefficients = context.initial_params.color_coefficients
color_coefficients[:n_params_color] = theta[start:end] * optimization_manager.weight_color_params
i = n_params_shape + n_params_expression + n_params_color

if optimization_manager.fix_camera_pose:
camera_pose = None
Expand All @@ -820,24 +885,28 @@ def from_theta(optimization_manager: Union[BFMOptimization, BFMOptimizationConte
optimization_manager=optimization_manager,
shape_coefficients=shape_coefficients,
expression_coefficients=expression_coefficients,
color_coefficients=color_coefficients,
camera_pose=camera_pose)

def with_new_manager(self, optimization_manager: BFMOptimization):
return BFMOptimizationParameters(
optimization_manager=optimization_manager,
shape_coefficients=self.shape_coefficients,
expression_coefficients=self.expression_coefficients,
color_coefficients=self.color_coefficients,
camera_pose=self.camera_pose
)

def to_theta(self):
theta = []
# To translate the parameters back into a theta list, shape and expression coefficients have to be divided
# again by their weights
# To translate the parameters back into a theta list, shape,
# expression & color coefficients have to be divided again by their weights
theta.extend(self.shape_coefficients[:self.optimization_manager.n_params_shape]
/ self.optimization_manager.weight_shape_params)
theta.extend(self.expression_coefficients[:self.optimization_manager.n_params_expression]
/ self.optimization_manager.weight_expression_params)
theta.extend(self.color_coefficients[:self.optimization_manager.n_params_color]
/ self.optimization_manager.weight_color_params)

if not self.optimization_manager.fix_camera_pose:
mode = self.optimization_manager.rotation_mode
Expand Down
Loading