diffrp.utils.shader_ops#

GLSL Attribute Access#

In addition to the utility methods, the module also contain a syntactic sugar to access channels of the last dimension in tensors conveniently, in a GLSL-like manner.

This would be particularly useful for graphics applications.

tnsr.xyzw can access the [0, 1, 2, 3] elements of the last channel, respectively. The same goes with tnsr.rgba.

You can combine any of them according to your need like a.xx, b.wxy, c.rgb, d.a, etc..

The singleton dimensions are not removed. If a has shape (N, 3), a.x will have shape (N, 1).

diffrp.utils.shader_ops.black_tex()#

Get a full black (zeros) RGBA image of shape [16, 16, 4].

diffrp.utils.shader_ops.blur2d(texture2d: Tensor, size: int, iters: int)#

Blur an 2D image (HWC) by the box kernel for several iterations, efficiently simulating a Gaussian-like blur.

Parameters:

  • texture2d – HWC tensor to blur.

  • size – Size of box kernel.

  • iters – Iterations of box blur.

Returns:

Has the same shape as texture2d.

Return type:

torch.Tensor

diffrp.utils.shader_ops.cross(a: Tensor, b: Tensor)#
diffrp.utils.shader_ops.dot(a: Tensor, b: Tensor, keepdim: bool = True)#
diffrp.utils.shader_ops.empty_normal_tex()#

Get a full empty ([0.5, 0.5, 1.0]) normal image of shape [16, 16, 3].

diffrp.utils.shader_ops.flipper_2d()#
diffrp.utils.shader_ops.flipper_3d()#
diffrp.utils.shader_ops.float2(*tensors)#

See floatx. The result should have 2 as the last channel dimension.

diffrp.utils.shader_ops.float3(*tensors)#

See floatx. The result should have 3 as the last channel dimension.

diffrp.utils.shader_ops.float4(*tensors)#

See floatx. The result should have 4 as the last channel dimension.

diffrp.utils.shader_ops.floatx(*tensors)#

Make a new tensor from sub-tensors or floats.

Inputs are broadcast and concatenated on the last dimension.

Example: floatx(rgba.rgb, 1) can give a new RGBA tensor with alpha set to opaque.

diffrp.utils.shader_ops.fma(x: Tensor, y: Tensor, b: Tensor, a: float = 1.0)#

a * x * y + b (multiply-add), but fused/faster.

diffrp.utils.shader_ops.from_c1hw(c1hw: Tensor)#

Convert a C1HW tensor back to HWC format.

diffrp.utils.shader_ops.fsa(a: float, x: Tensor, b: Tensor | float)#

a * x + b (scale-add), but fused/faster

diffrp.utils.shader_ops.full_like_vec(x: Tensor, value, v)#

Fill in value like the given tensor, but change the last shape dimension to v.

diffrp.utils.shader_ops.gpu_color(inputs)#

Cast any input (float, list, numpy array, tensor) into torch.Tensor with dtype float32 on the current CUDA device. If any of the inputs have a maximum value greater than 1.5, we assume the input to be uint8 coded and divide the colors by 255. CUDA_VISIBLE_DEVICES and torch.cuda.set_device can be used to specify the current device.

diffrp.utils.shader_ops.gpu_f32(inputs)#

Cast any input (float, list, numpy array, tensor) into torch.Tensor with dtype float32 on the current CUDA device. CUDA_VISIBLE_DEVICES and torch.cuda.set_device can be used to specify the current device.

diffrp.utils.shader_ops.gpu_i32(inputs)#

Cast any input (float, list, numpy array, tensor) into torch.Tensor with dtype int32 on the current CUDA device. CUDA_VISIBLE_DEVICES and torch.cuda.set_device can be used to specify the current device.

diffrp.utils.shader_ops.gray_tex()#

Get a full gray (0.5) RGBA image of shape [16, 16, 4].

diffrp.utils.shader_ops.homogeneous(coords: Tensor)#

Concatenate a channel of ones to the given coords to obtain its homogeneous-coordinate position.

diffrp.utils.shader_ops.homogeneous_vec(coords: Tensor)#

Concatenate a channel of zeros to the given coords to obtain its homogeneous-coordinate direction.

diffrp.utils.shader_ops.hsv2rgb(hsv: Tensor) → Tensor#

Convert HSV to RGB. Broadcastable.

diffrp.utils.shader_ops.length(x: Tensor)#

Compute the length of x vector over the last dimension.

The method is broadcastable.

Parameters:

x (torch.Tensor) – Shape (…, C).

Returns:

Shape (…, 1), L2 length of vectors.

Return type:

torch.Tensor

diffrp.utils.shader_ops.multi_surf(*surfs)#
diffrp.utils.shader_ops.normalized(x: Tensor)#

Normalize x vector over the last dimension. Zero vectors are kept as-is.

The method is broadcastable.

diffrp.utils.shader_ops.ones_like_vec(x: Tensor, v)#

Ones like the given tensor, but change the last shape dimension to v.

diffrp.utils.shader_ops.reflect(i: Tensor, n: Tensor)#

Reflects a ray according to a normal vector.

The method is broadcastable.

Parameters:

  • i – Inbound ray. Shape (…, 3).

  • n – Normal. Shape (…, 3).

Returns:

Tensor of the same shape as inputs.

Return type:

torch.Tensor

diffrp.utils.shader_ops.rgb2hsv(rgb: Tensor) → Tensor#

Convert RGB to HSV. Broadcastable.

diffrp.utils.shader_ops.sample2d(texture2d: Tensor, texcoords: Tensor, wrap: str = 'border', mode: str = 'bilinear')#

Samples a 2D texture on UV coords.

Parameters:

  • texture2d (torch.Tensor) – Tensor of shape (H, W, C), the texture to be sampled.

  • texcoords (torch.Tensor) – Tensor of shape (…, 2), batch of coordinates to sample. (0, 0) is the bottom-left corner (last row in H axis).

  • wrap (str) – Defines how the boundaries and points outside boundaries are handled. One of “border”, “reflection” and “cyclic”.

  • mode (str) – Sampling method. One of “nearest”, “bilinear” and “bicubic”.

Returns:

Tensor of shape (…, C), sampled texture.

The batch dimensions are kept the same as texcoords.

Return type:

torch.Tensor

diffrp.utils.shader_ops.sample3d(texture3d: Tensor, texcoords: Tensor, wrap: str = 'border', mode: str = 'bilinear')#

Samples a 3D texture on UV coords.

Parameters:

  • texture3d (torch.Tensor) – Tensor of shape (D, H, W, C), the texture to be sampled.

  • texcoords (torch.Tensor) – Tensor of shape (…, 3), batch of coordinates to sample. (0, 0, 0) is the bottom-left corner, and the first slice of the depth.

  • wrap (str) – Defines how the boundaries and points outside boundaries are handled. One of “border” and “reflection”.

  • mode (str) – Sampling method. One of “nearest” and “bilinear”.

Returns:

Tensor of shape (…, C), sampled texture.

The batch dimensions are kept the same as texcoords.

Return type:

torch.Tensor

diffrp.utils.shader_ops.saturate(x: Tensor)#

Clamps values in x to be between 0 and 1.

The method is element-wise.

diffrp.utils.shader_ops.small_matrix_inverse(x: Tensor) → Tensor#

Faster inverse for matrices smaller than 2x8x8.

diffrp.utils.shader_ops.sobel(texture2d: Tensor)#

Compute the Sobel edge detection operator on a grayscale image.

Parameters:

texture2d (torch.Tensor) – (H, W, 1) grayscale image.

Returns:

(H, W, 4) containing (strength_x, strength_y, strength_total, input).

Return type:

torch.Tensor

diffrp.utils.shader_ops.split_alpha(x: Tensor)#

Splits x on the last dimension. The singleton alpha dimension is kept.

The method is broadcastable.

Parameters:

x (torch.Tensor) – Shape (…, C + 1).

Returns:

Shapes (…, C) and (…, 1), respectively.

Return type:

Tuple of tensors (c, alpha)

diffrp.utils.shader_ops.to_bchw(hwc: Tensor)#

Convert a HWC or BHWC tensor into BCHW format.

diffrp.utils.shader_ops.to_c1hw(hwc: Tensor)#

Convert a HWC tensor into C1HW format, i.e. a batch of single-channel images in BCHW format.

diffrp.utils.shader_ops.to_hwc(bchw: Tensor)#

Convert a BCHW or CHW tensor back to HWC format.

diffrp.utils.shader_ops.transform_point(xyz, matrix)#

Transform points xyz with the homogeneous transform matrix.

The last dimension of xyz needs to be 3. The result will have the same shape as xyz but the last dimension changed to 4.

matrix is a tensor of shape (4, 4) and the transform is applied as a left-multiplication to xyz.

The method is fully broadcastable to any dimensions.

diffrp.utils.shader_ops.transform_point4x3(xyz, matrix)#

Transform points xyz with the homogeneous transform matrix. The matrix needs to be affine. The last dimension of xyz needs to be 3.

This is in general faster than doing a transform_point and then perform a homogeneous division.

The result will have the same shape as xyz, unlike transform_point which keeps the homogeneous coordinate form.

xyz should be of shape (B, 3).

diffrp.utils.shader_ops.transform_vector(xyz, matrix)#

Transform a vector (direction) xyz by the homogeneous transform matrix. The last dimension of xyz needs to be 3.

The result will have the same shape as xyz.

matrix is a tensor of shape (4, 4) and the transform is applied as a left-multiplication to xyz.

The method is fully broadcastable to any dimensions.

diffrp.utils.shader_ops.transform_vector3x3(xyz, matrix)#

Transform a vector (direction) xyz by the homogeneous transform matrix. The last dimension of xyz needs to be 3.

The matrix needs to be affine.

The result will have the same shape as xyz.

matrix is a tensor of shape (4, 4) and the transform is applied as a left-multiplication to xyz.

The method is fully broadcastable to any dimensions.

diffrp.utils.shader_ops.white_tex()#

Get a full white (ones) RGBA image of shape [16, 16, 4].

diffrp.utils.shader_ops.zeros_like_vec(x: Tensor, v)#

Zeros like the given tensor, but change the last shape dimension to v.