# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# pyre-strict
from typing import Callable, Optional, Tuple
import torch as th
import torch.nn.functional as thf
from drtk.interpolate import interpolate
from drtk.utils import index, load_torch_ops
load_torch_ops("drtk.edge_grad_ext")
[docs]
@th.compiler.disable
def edge_grad_estimator(
v_pix: th.Tensor,
vi: th.Tensor,
bary_img: th.Tensor,
img: th.Tensor,
index_img: th.Tensor,
v_pix_img_hook: Optional[Callable[[th.Tensor], None]] = None,
max_dp_dr: float = 1e4,
) -> th.Tensor:
"""Makes the rasterized image ``img`` differentiable at visibility discontinuities
and backpropagates the gradients to ``v_pix``.
This function takes a rasterized image ``img`` that is assumed to be differentiable at
continuous regions but not at discontinuities. In some cases, ``img`` may not be differentiable
at all. For example, if the image is a rendered segmentation mask, it remains constant at
continuous regions, making it non-differentiable. However, ``edge_grad_estimator`` can still
compute gradients at the discontinuities with respect to ``v_pix``.
The arguments ``bary_img`` and ``index_img`` must correspond exactly to the rasterized image
``img``. Each pixel in ``img`` should correspond to a fragment originated prom primitive
specified by ``index_img`` and it should have barycentric coordinates specified by
``bary_img``. This means that with a small change to ``v_pix``, the pixels in ``img`` should
change accordingly. A frequent mistake that violates this condition is applying a mask
to the rendered image to exclude unwanted regions, which leads to erroneous gradients.
The function returns the ``img`` unchanged but with added differentiability at the
discontinuities. Note that it is not necessary for the input ``img`` to require gradients,
but the returned ``img`` will always require gradients.
Args:
v_pix (Tensor): Pixel-space vertex coordinates, preserving the original camera-space
Z-values. Shape: :math:`(N, V, 3)`.
vi (Tensor): Face vertex index list tensor. Shape: :math:`(F, 3)` or :math:`(N, F, 3)`.
bary_img (Tensor): 3D barycentric coordinate image tensor. Shape: :math:`(N, 3, H, W)`.
img (Tensor): The rendered image. Shape: :math:`(N, C, H, W)`.
index_img (Tensor): Index image tensor. Shape: :math:`(N, H, W)`.
v_pix_img_hook (Optional[Callable[[th.Tensor], None]]): An optional backward hook that will
be registered to ``v_pix_img``. Useful for examining the generated image space. Default
is None.
max_dp_dr (float): Maximum allowed magnitude for the edge-to-fragment
derivative ∂p/∂r (Eqn. 14 of the paper). Prevents gradient blowup
for nearly-coplanar triangles by clamping ‖∂p/∂r‖ ≤ max_dp_dr.
Set to 0.0 to disable clamping (for comparison with analytic
solutions or finite differences). Default is 1e4.
**How to choose max_dp_dr:** Because ``v_pix`` has x,y in pixel
units but z in world (camera-space) units, the normals computed
from ``v_pix`` are in a mixed-unit space. This means the effective
world-space dihedral angle below which clamping activates depends
on the camera's focal length and the distance to the geometry::
θ ≈ f / (max_dp_dr · Z)
where ``f`` is the focal length in pixels and ``Z`` is the depth
in world units. Examples with default max_dp_dr=1e4:
========== ====== ==============
f (pixels) Z θ (degrees)
========== ====== ==============
500 5 0.6°
1111 4 1.6°
5000 1 29°
10000 1000 0.06°
========== ====== ==============
For typical rendering setups (moderate focal length, geometry not
too close), the default 1e4 clips only at sub-degree to few-degree
dihedral angles. For very high ``f/Z`` ratios (telephoto lens or
close-up geometry), consider increasing max_dp_dr.
Ideally, the clamping would operate on normals in world space
(where the dihedral angle is independent of camera parameters).
This would require passing the focal length to the kernel and
rescaling the z-component of the projected normals to pixel units
before computing ∂p/∂r. This is not yet implemented — the current
approach operates directly on the mixed-unit ``v_pix`` normals.
Returns:
Tensor: Returns the input ``img`` unchanged. However, the returned image now has added
differentiability at visibility discontinuities. This returned image should be used for
computing losses
Note:
It is crucial not to spatially modify the rasterized image before passing it to
`edge_grad_estimator`. That stems from the requirement that ``bary_img`` and ``index_img``
must correspond exactly to the rasterized image ``img``. That means that the location of all
discontinuities is controlled by ``v_pix`` and can be modified by modifing ``v_pix``.
Operations that are allowed, as long as they are differentiable, include:
- Pixel-wise MLP
- Color mapping
- Color correction, gamma correction
- Anything that would be indistinguishable from processing fragments independently
before their values get assigned to pixels of ``img``
Operations that **must be avoided** before `edge_grad_estimator` include:
- Gaussian blur
- Warping or deformation
- Masking, cropping, or introducing holes
There is however, no issue with appling them after `edge_grad_estimator`.
If the operation is highly non-linear, it is recommended to perform it before calling
:func:`edge_grad_estimator`.
All sorts of clipping and clamping (e.g., `x.clamp(min=0.0, max=1.0)`) must also be done
before invoking this function.
Usage Example::
import torch.nn.functional as thf
from drtk import transform, rasterize, render, interpolate, edge_grad_estimator
...
v_pix = transform(v, tex, campos, camrot, focal, princpt)
index_img = rasterize(v_pix, vi, width=512, height=512)
_, bary_img = render(v_pix, vi, index_img)
vt_img = interpolate(vt, vti, index_img, bary_img)
img = thf.grid_sample(
tex,
vt_img.permute(0, 2, 3, 1),
mode="bilinear",
padding_mode="border",
align_corners=False
)
mask = (index_img != -1)[:, None, :, :]
img = img * mask
img = edge_grad_estimator(
v_pix=v_pix,
vi=vi,
bary_img=bary_img,
img=img,
index_img=index_img
)
optim.zero_grad()
image_loss = loss_func(img, img_gt)
image_loss.backward()
optim.step()
"""
if vi.ndim == 2:
vi = vi[None, ...].expand(v_pix.shape[0], -1, -1)
# TODO: avoid call to interpolate, use backward kernel of interpolate directly
# Doing so will make `edge_grad_estimator` zero-overhead in forward pass
# At the moment, value of `v_pix_img` is ignored, and only passed to
# edge_grad_estimator so that backward kernel can be called with the computed gradient.
v_pix_img = interpolate(v_pix, vi, index_img, bary_img.detach())
img = th.ops.edge_grad_ext.edge_grad_estimator(
v_pix, v_pix_img, vi, img, index_img, max_dp_dr
)
if v_pix_img_hook is not None:
v_pix_img.register_hook(v_pix_img_hook)
return img
[docs]
def edge_grad_estimator_ref(
v_pix: th.Tensor,
vi: th.Tensor,
bary_img: th.Tensor,
img: th.Tensor,
index_img: th.Tensor,
v_pix_img_hook: Optional[Callable[[th.Tensor], None]] = None,
max_dp_dr: float = 1e4,
) -> th.Tensor:
"""Pure PyTorch reference implementation used by tests.
This helper is intentionally not part of the documented public API. See
:func:`drtk.edge_grad_estimator` for the supported implementation.
"""
# could use v_pix_img output from DRTK, but bary_img needs to be detached.
v_pix_img = interpolate(v_pix, vi, index_img, bary_img.detach())
# pyre-fixme[16]: `EdgeGradEstimatorFunction` has no attribute `apply`.
img = EdgeGradEstimatorFunction.apply(
v_pix, v_pix_img, vi, img, index_img, max_dp_dr
)
if v_pix_img_hook is not None:
v_pix_img.register_hook(v_pix_img_hook)
return img
class EdgeGradEstimatorFunction(th.autograd.Function):
@staticmethod
# pyre-fixme[14]: `forward` overrides method defined in `Function` inconsistently.
def forward(
ctx,
v_pix: th.Tensor,
v_pix_img: th.Tensor,
vi: th.Tensor,
img: th.Tensor,
index_img: th.Tensor,
max_dp_dr: float = 1e4,
) -> th.Tensor:
ctx.save_for_backward(v_pix, img, index_img, vi)
ctx.max_dp_dr = max_dp_dr
return img
@staticmethod
# pyre-fixme[14]: `backward` overrides method defined in `Function` inconsistently.
def backward(
ctx, grad_output: th.Tensor
) -> Tuple[
Optional[th.Tensor],
Optional[th.Tensor],
Optional[th.Tensor],
Optional[th.Tensor],
Optional[th.Tensor],
None,
]:
# early exit in case geometry is not optimized.
if not ctx.needs_input_grad[1]:
return None, None, None, grad_output, None, None
max_dp_dr = ctx.max_dp_dr
v_pix, img, index_img, vi = ctx.saved_tensors
x_grad = img[:, :, :, 1:] - img[:, :, :, :-1]
y_grad = img[:, :, 1:, :] - img[:, :, :-1, :]
l_index = index_img[:, None, :, :-1]
r_index = index_img[:, None, :, 1:]
t_index = index_img[:, None, :-1, :]
b_index = index_img[:, None, 1:, :]
x_mask = r_index != l_index
y_mask = b_index != t_index
x_both_triangles = (r_index != -1) & (l_index != -1)
y_both_triangles = (b_index != -1) & (t_index != -1)
iimg_clamped = index_img.clamp(min=0).long()
# compute barycentric coordinates
b = v_pix.shape[0]
vi_img = index(vi, iimg_clamped, 0).long()
p0 = th.cat(
[index(v_pix[i], vi_img[i, ..., 0].data, 0)[None, ...] for i in range(b)],
dim=0,
)
p1 = th.cat(
[index(v_pix[i], vi_img[i, ..., 1].data, 0)[None, ...] for i in range(b)],
dim=0,
)
p2 = th.cat(
[index(v_pix[i], vi_img[i, ..., 2].data, 0)[None, ...] for i in range(b)],
dim=0,
)
v10 = p1 - p0
v02 = p0 - p2
n = th.cross(v02, v10)
px, py = th.meshgrid(
th.arange(img.shape[-2], device=v_pix.device),
th.arange(img.shape[-1], device=v_pix.device),
)
def epsclamp(x: th.Tensor) -> th.Tensor:
return th.where(x < 0, x.clamp(max=-1e-8), x.clamp(min=1e-8))
# pyre-fixme[53]: Captured variable `n` is not annotated.
# pyre-fixme[53]: Captured variable `p0` is not annotated.
# pyre-fixme[53]: Captured variable `px` is not annotated.
# pyre-fixme[53]: Captured variable `py` is not annotated.
# pyre-fixme[53]: Captured variable `v02` is not annotated.
# pyre-fixme[53]: Captured variable `v10` is not annotated.
def check_if_point_inside_triangle(offset_x: int, offset_y: int) -> th.Tensor:
_px = px + offset_x
_py = py + offset_y
vp0p = th.stack([p0[..., 0] - _px, p0[..., 1] - _py], dim=-1) / epsclamp(
n[..., 2:3]
)
bary_1 = v02[..., 0] * -vp0p[..., 1] + v02[..., 1] * vp0p[..., 0]
bary_2 = v10[..., 0] * -vp0p[..., 1] + v10[..., 1] * vp0p[..., 0]
return ((bary_1 > 0) & (bary_2 > 0) & ((bary_1 + bary_2) < 1))[:, None]
left_pnt_inside_right_triangle = (
check_if_point_inside_triangle(-1, 0)[..., :, 1:]
& x_mask
& x_both_triangles
)
right_pnt_inside_left_triangle = (
check_if_point_inside_triangle(1, 0)[..., :, :-1]
& x_mask
& x_both_triangles
)
down_pnt_inside_up_triangle = (
check_if_point_inside_triangle(0, 1)[..., :-1, :]
& y_mask
& y_both_triangles
)
up_pnt_inside_down_triangle = (
check_if_point_inside_triangle(0, -1)[..., 1:, :]
& y_mask
& y_both_triangles
)
horizontal_intersection = (
right_pnt_inside_left_triangle & left_pnt_inside_right_triangle
)
vertical_intersection = (
down_pnt_inside_up_triangle & up_pnt_inside_down_triangle
)
left_hangs_over_right = left_pnt_inside_right_triangle & (
~right_pnt_inside_left_triangle
)
right_hangs_over_left = right_pnt_inside_left_triangle & (
~left_pnt_inside_right_triangle
)
up_hangs_over_down = up_pnt_inside_down_triangle & (
~down_pnt_inside_up_triangle
)
down_hangs_over_up = down_pnt_inside_up_triangle & (
~up_pnt_inside_down_triangle
)
x_grad *= x_mask
y_grad *= y_mask
grad_output_x = 0.5 * (grad_output[:, :, :, 1:] + grad_output[:, :, :, :-1])
grad_output_y = 0.5 * (grad_output[:, :, 1:, :] + grad_output[:, :, :-1, :])
x_grad = (x_grad * grad_output_x).sum(dim=1)
y_grad = (y_grad * grad_output_y).sum(dim=1)
x_grad_no_int = x_grad * (~horizontal_intersection[:, 0])
y_grad_no_int = y_grad * (~vertical_intersection[:, 0])
x_grad_spread = th.zeros(
*x_grad_no_int.shape[:1],
x_grad_no_int.shape[1],
y_grad_no_int.shape[2],
dtype=x_grad_no_int.dtype,
device=x_grad_no_int.device,
)
x_grad_spread[:, :, :-1] = x_grad_no_int * (~right_hangs_over_left[:, 0])
x_grad_spread[:, :, 1:] += x_grad_no_int * (~left_hangs_over_right[:, 0])
y_grad_spread = th.zeros_like(x_grad_spread)
y_grad_spread[:, :-1, :] = y_grad_no_int * (~down_hangs_over_up[:, 0])
y_grad_spread[:, 1:, :] += y_grad_no_int * (~up_hangs_over_down[:, 0])
# Intersections. Compute border sliding gradients
#################################################
z_grad_spread = th.zeros_like(x_grad_spread)
x_grad_int = x_grad * horizontal_intersection[:, 0]
y_grad_int = y_grad * vertical_intersection[:, 0]
n = thf.normalize(n, dim=-1)
n = n.permute(0, 3, 1, 2)
n_left = n[..., :, :-1]
n_right = n[..., :, 1:]
n_up = n[..., :-1, :]
n_down = n[..., 1:, :]
def get_dp_dr(v_varying: th.Tensor, v_fixed: th.Tensor) -> th.Tensor:
"""
Computes ∂p/∂r — the derivative of the edge position p with
respect to fragment position r, projected onto a 2D plane (XZ
for vertical edges, YZ for horizontal edges). See Eqn. 14 of
the paper.
When max_dp_dr > 0, the result magnitude is clamped to at most
max_dp_dr to stay within the linear model's domain of validity.
See drtk/src/edge_grad/edge_grad_kernel.cu get_dp_dr() for the
full derivation.
"""
v_varying = thf.normalize(v_varying, dim=1)
v_fixed = thf.normalize(v_fixed, dim=1)
b = th.stack([-v_fixed[:, 1], v_fixed[:, 0]], dim=1)
d = (b * v_varying).sum(dim=1, keepdim=True)
if max_dp_dr > 0:
# Clamp |b.x/d| ≤ M by ensuring |d| ≥ |b.x|/M.
# Use where() instead of sign() because sign(0)=0 would zero out safe_d.
sign_d = th.where(d >= 0, th.ones_like(d), -th.ones_like(d))
safe_d = sign_d * th.maximum(d.abs(), b[:, 0:1].abs() / max_dp_dr)
# safe_d can only be zero if both d and b.x are zero, which means
# n_fixed is along the first axis — epsclamp handles this edge case.
return b[:, 0:1] / epsclamp(safe_d) * v_varying
else:
return b[:, 0:1] / epsclamp(d) * v_varying
# We compute partial derivatives by fixing one triangle and moving the
# other, and then vice versa.
# Left triangle moves, right fixed
dpx_dr = get_dp_dr(n_left[:, [0, 2]], -n_right[:, [0, 2]])
x_grad_spread[:, :, :-1] += x_grad_int * dpx_dr[:, 0]
z_grad_spread[:, :, :-1] += x_grad_int * dpx_dr[:, 1]
# Left triangle fixed, right moves
dpx_dr = get_dp_dr(n_right[:, [0, 2]], n_left[:, [0, 2]])
x_grad_spread[:, :, 1:] += x_grad_int * dpx_dr[:, 0]
z_grad_spread[:, :, 1:] += x_grad_int * dpx_dr[:, 1]
# Upper triangle moves, lower fixed
dpy_dr = get_dp_dr(n_up[:, [1, 2]], -n_down[:, [1, 2]])
y_grad_spread[:, :-1, :] += y_grad_int * dpy_dr[:, 0]
z_grad_spread[:, :-1, :] += y_grad_int * dpy_dr[:, 1]
# Lower triangle moves, upper fixed
dpy_dr = get_dp_dr(n_down[:, [1, 2]], n_up[:, [1, 2]])
y_grad_spread[:, 1:, :] += y_grad_int * dpy_dr[:, 0]
z_grad_spread[:, 1:, :] += y_grad_int * dpy_dr[:, 1]
m = index_img == -1
x_grad_spread[m] = 0.0
y_grad_spread[m] = 0.0
grad_v_pix = -th.stack([x_grad_spread, y_grad_spread, z_grad_spread], dim=3)
return None, grad_v_pix, None, grad_output, None, None