drtk.edge_grad_estimator#

`edge_grad_estimator`	Makes the rasterized image `img` differentiable at visibility discontinuities and backpropagates the gradients to `v_pix`.
`edge_grad_estimator_ref`	Pure PyTorch reference implementation used by tests.

drtk.edge_grad_estimator(v_pix, vi, bary_img, img, index_img, v_pix_img_hook=None, max_dp_dr=10000.0)[source]#

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.

Parameters:

v_pix (Tensor) – Pixel-space vertex coordinates, preserving the original camera-space Z-values. Shape: \((N, V, 3)\).
vi (Tensor) – Face vertex index list tensor. Shape: \((F, 3)\) or \((N, F, 3)\).
bary_img (Tensor) – 3D barycentric coordinate image tensor. Shape: \((N, 3, H, W)\).
img (Tensor) – The rendered image. Shape: \((N, C, H, W)\).
index_img (Tensor) – Index image tensor. Shape: \((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:

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

Return type:

Tensor

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 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()

drtk.edge_grad_estimator_ref(v_pix, vi, bary_img, img, index_img, v_pix_img_hook=None, max_dp_dr=10000.0)[source]#

Pure PyTorch reference implementation used by tests.

This helper is intentionally not part of the documented public API. See drtk.edge_grad_estimator() for the supported implementation.

f (pixels)	Z	θ (degrees)
500	5	0.6°
1111	4	1.6°
5000	1	29°
10000	1000	0.06°