# # The original code is under the following copyright: # Copyright (C) 2023, Inria # GRAPHDECO research group, https://team.inria.fr/graphdeco # All rights reserved. # # This software is free for non-commercial, research and evaluation use # under the terms of the LICENSE_GS.md file. # # For inquiries contact george.drettakis@inria.fr # # The modifications of the code are under the following copyright: # Copyright (C) 2024, University of Liege, KAUST and University of Oxford # TELIM research group, http://www.telecom.ulg.ac.be/ # IVUL research group, https://ivul.kaust.edu.sa/ # VGG research group, https://www.robots.ox.ac.uk/~vgg/ # All rights reserved. # The modifications are under the LICENSE.md file. # # For inquiries contact jan.held@uliege.be # from typing import NamedTuple import torch.nn as nn import torch from . import _C def cpu_deep_copy_tuple(input_tuple): copied_tensors = [item.cpu().clone() if isinstance(item, torch.Tensor) else item for item in input_tuple] return tuple(copied_tensors) def rasterize_triangles( triangles_points, sigma, num_points_per_triangle, cumsum_of_points_per_triangle, number_of_points, sh, colors_precomp, opacities, means2D, scaling, density_factor, raster_settings, ): return _RasterizeTriangles.apply( triangles_points, sigma, num_points_per_triangle, cumsum_of_points_per_triangle, number_of_points, sh, colors_precomp, opacities, means2D, scaling, density_factor, raster_settings, ) class _RasterizeTriangles(torch.autograd.Function): @staticmethod def forward( ctx, triangles_points, sigma, num_points_per_triangle, cumsum_of_points_per_triangle, number_of_points, sh, colors_precomp, opacities, means2D, scaling, density_factor, raster_settings, ): # Restructure arguments the way that the C++ lib expects them args = ( raster_settings.bg, triangles_points, sigma, num_points_per_triangle, cumsum_of_points_per_triangle, colors_precomp, opacities, scaling, density_factor, raster_settings.viewmatrix, raster_settings.projmatrix, number_of_points, raster_settings.tanfovx, raster_settings.tanfovy, raster_settings.image_height, raster_settings.image_width, sh, raster_settings.sh_degree, raster_settings.campos, raster_settings.prefiltered, raster_settings.debug ) # Invoke C++/CUDA rasterizer if raster_settings.debug: cpu_args = cpu_deep_copy_tuple(args) # Copy them before they can be corrupted try: num_rendered, color, depth, radii, geomBuffer, binningBuffer, imgBuffer, scaling, density_factor, max_blending = _C.rasterize_triangles(*args) except Exception as ex: torch.save(cpu_args, "snapshot_fw.dump") print("\nAn error occured in forward. Please forward snapshot_fw.dump for debugging.") raise ex else: num_rendered, color, depth, radii, geomBuffer, binningBuffer, imgBuffer, scaling, density_factor, max_blending = _C.rasterize_triangles(*args) # Keep relevant tensors for backward ctx.raster_settings = raster_settings ctx.num_rendered = num_rendered ctx.number_of_points = number_of_points ctx.save_for_backward(triangles_points, sigma, num_points_per_triangle, cumsum_of_points_per_triangle, colors_precomp, radii, sh, geomBuffer, binningBuffer, imgBuffer) return color, radii, scaling, density_factor, depth, max_blending @staticmethod def backward(ctx, grad_out_color, _, __, ___, grad_depth, _____): # Restore necessary values from context num_rendered = ctx.num_rendered raster_settings = ctx.raster_settings number_of_points = ctx.number_of_points triangles_points, sigma, num_points_per_triangle, cumsum_of_points_per_triangle, colors_precomp, radii, sh, geomBuffer, binningBuffer, imgBuffer = ctx.saved_tensors # Restructure args as C++ method expects them args = (raster_settings.bg, triangles_points, sigma, num_points_per_triangle, cumsum_of_points_per_triangle, radii, colors_precomp, raster_settings.viewmatrix, raster_settings.projmatrix, number_of_points, raster_settings.tanfovx, raster_settings.tanfovy, grad_out_color, grad_depth, sh, raster_settings.sh_degree, raster_settings.campos, geomBuffer, num_rendered, binningBuffer, imgBuffer, raster_settings.debug) # Compute gradients for relevant tensors by invoking backward method if raster_settings.debug: cpu_args = cpu_deep_copy_tuple(args) # Copy them before they can be corrupted try: grad_triangles, grad_sigma, grad_colors_precomp, grad_opacities, grad_sh, grad_means2D = _C.rasterize_triangles_backward(*args) except Exception as ex: torch.save(cpu_args, "snapshot_bw.dump") print("\nAn error occured in backward. Writing snapshot_bw.dump for debugging.\n") raise ex else: grad_triangles, grad_sigma, grad_colors_precomp, grad_opacities, grad_sh, grad_means2D = _C.rasterize_triangles_backward(*args) #print(torch.max(torch.abs(grad_triangles)), torch.min(torch.abs(grad_triangles))) #grad_triangles = grad_triangles.reshape(-1, 8, 3) grad_triangles = grad_triangles.flatten(0) grad_sigma = grad_sigma.view(-1, 1) grads = ( grad_triangles, grad_sigma, None, None, None, grad_sh, grad_colors_precomp, grad_opacities, grad_means2D, None, None, None ) return grads class TriangleRasterizationSettings(NamedTuple): image_height: int image_width: int tanfovx : float tanfovy : float bg : torch.Tensor scale_modifier : float viewmatrix : torch.Tensor projmatrix : torch.Tensor sh_degree : int campos : torch.Tensor prefiltered : bool debug : bool class TriangleRasterizer(nn.Module): def __init__(self, raster_settings): super().__init__() self.raster_settings = raster_settings def markVisible(self, positions): # Mark visible points (based on frustum culling for camera) with a boolean with torch.no_grad(): raster_settings = self.raster_settings visible = _C.mark_visible( positions, raster_settings.viewmatrix, raster_settings.projmatrix) return visible def forward(self, triangles_points, sigma, num_points_per_triangle, cumsum_of_points_per_triangle, number_of_points, opacities, means2D, scaling, density_factor, shs = None, colors_precomp = None): raster_settings = self.raster_settings if (shs is None and colors_precomp is None) or (shs is not None and colors_precomp is not None): raise Exception('Please provide excatly one of either SHs or precomputed colors!') if shs is None: shs = torch.Tensor([]) if colors_precomp is None: colors_precomp = torch.Tensor([]) # Invoke C++/CUDA rasterization routine return rasterize_triangles( triangles_points, sigma, num_points_per_triangle, cumsum_of_points_per_triangle, number_of_points, shs, colors_precomp, opacities, means2D, scaling, density_factor, raster_settings, )