11.1 C
New York

A Python Engineer’s Introduction to 3D Gaussian Splatting (Part 3) | by Derek Austin | Jul, 2024

Part 3 of our Gaussian Splatting tutorial, showing how to render splats onto a 2D imageFinally, we reach the most intriguing phase of the Gaussian splatting process: rendering! This step is arguably the most crucial, as it determines the realism of our model. Yet, it might also be the simplest. In part 1 and part 2 of our series we demonstrated how to transform raw splats into a format ready for rendering, but now we actually have to do the work and render onto a fixed set of pixels. The authors have developed a fast rendering engine using CUDA, which can be somewhat challenging to follow. Therefore, I believe it is beneficial to first walk through the code in Python, using straightforward for loops for clarity. For those eager to dive deeper, all the necessary code is available on our GitHub.Let’s discuss how to render each individual pixel. From our previous article, we have all the necessary components: 2D points, associated colors, covariance, sorted depth order, inverse covariance in 2D, minimum and maximum x and y values for each splat, and associated opacity. With these components, we can render any pixel. Given specific pixel coordinates, we iterate through all splats until we reach a saturation threshold, following the splat depth order relative to the camera plane (projected to the camera plane and then sorted by depth). For each splat, we first check if the pixel coordinate is within the bounds defined by the minimum and maximum x and y values. This check determines if we should continue rendering or ignore the splat for these coordinates. Next, we compute the Gaussian splat strength at the pixel coordinate using the splat mean, splat covariance, and pixel coordinates.def compute_gaussian_weight(pixel_coord: torch.Tensor, # (1, 2) tensorpoint_mean: torch.Tensor,inverse_covariance: torch.Tensor,) -> torch.Tensor:difference = point_mean – pixel_coordpower = -0.5 * difference @ inverse_covariance @ difference.Treturn torch.exp(power).item()We multiply this weight by the splat’s opacity to obtain a parameter called alpha. Before adding this new value to the pixel, we need to check if we have exceeded our saturation threshold. We do not want a splat behind other splats to affect the pixel coloring and use computing resources if the pixel is already saturated. Thus, we use a threshold that allows us to stop rendering once it is exceeded. In practice, we start our saturation threshold at 1 and then multiply it by min(0.99, (1 — alpha)) to get a new value. If this value is less than our threshold (0.0001), we stop rendering that pixel and consider it complete. If not, we add the colors weighted by the saturation * (1 — alpha) value and update the saturation as new_saturation = old_saturation * (1 — alpha). Finally, we loop over every pixel (or every 16×16 tile in practice) and render. The complete code is shown below.def render_pixel(self,pixel_coords: torch.Tensor,points_in_tile_mean: torch.Tensor,colors: torch.Tensor,opacities: torch.Tensor,inverse_covariance: torch.Tensor,min_weight: float = 0.000001,) -> torch.Tensor:total_weight = torch.ones(1).to(points_in_tile_mean.device)pixel_color = torch.zeros((1, 1, 3)).to(points_in_tile_mean.device)for point_idx in range(points_in_tile_mean.shape[0]):point = points_in_tile_mean[point_idx, :].view(1, 2)weight = compute_gaussian_weight(pixel_coord=pixel_coords,point_mean=point,inverse_covariance=inverse_covariance[point_idx],)alpha = weight * torch.sigmoid(opacities[point_idx])test_weight = total_weight * (1 – alpha)if test_weight < min_weight:return pixel_colorpixel_color += total_weight * alpha * colors[point_idx]total_weight = test_weight# in case we never reach saturationreturn pixel_colorNow that we can render a pixel we can render a patch of an image, or what the authors refer to as a tile! def render_tile(self,x_min: int,y_min: int,points_in_tile_mean: torch.Tensor,colors: torch.Tensor,opacities: torch.Tensor,inverse_covariance: torch.Tensor,tile_size: int = 16,) -> torch.Tensor:”””Points in tile should be arranged in order of depth”””tile = torch.zeros((tile_size, tile_size, 3))# iterate by tiles for more efficient processingfor pixel_x in range(x_min, x_min + tile_size):for pixel_y in range(y_min, y_min + tile_size):tile[pixel_x % tile_size, pixel_y % tile_size] = self.render_pixel(pixel_coords=torch.Tensor([pixel_x, pixel_y]).view(1, 2).to(points_in_tile_mean.device),points_in_tile_mean=points_in_tile_mean,colors=colors,opacities=opacities,inverse_covariance=inverse_covariance,)return tileAnd finally we can use all of those tiles to render an entire image. Note how we check to make sure the splat will actually affect the current tile (x_in_tile and y_in_tile code).def render_image(self, image_idx: int, tile_size: int = 16) -> torch.Tensor:”””For each tile have to check if the point is in the tile”””preprocessed_scene = self.preprocess(image_idx)height = self.images[image_idx].heightwidth = self.images[image_idx].widthimage = torch.zeros((width, height, 3))for x_min in tqdm(range(0, width, tile_size)):x_in_tile = (x_min >= preprocessed_scene.min_x) & (x_min + tile_size <= preprocessed_scene.max_x)if x_in_tile.sum() == 0:continuefor y_min in range(0, height, tile_size):y_in_tile = (y_min >= preprocessed_scene.min_y) & (y_min + tile_size <= preprocessed_scene.max_y)points_in_tile = x_in_tile & y_in_tileif points_in_tile.sum() == 0:continuepoints_in_tile_mean = preprocessed_scene.points[points_in_tile]colors_in_tile = preprocessed_scene.colors[points_in_tile]opacities_in_tile = preprocessed_scene.sigmoid_opacity[points_in_tile]inverse_covariance_in_tile = preprocessed_scene.inverse_covariance_2d[points_in_tile]image[x_min : x_min + tile_size, y_min : y_min + tile_size] = (self.render_tile(x_min=x_min,y_min=y_min,points_in_tile_mean=points_in_tile_mean,colors=colors_in_tile,opacities=opacities_in_tile,inverse_covariance=inverse_covariance_in_tile,tile_size=tile_size,))return imageAt long last now that we have all the necessary components we can render an image. We take all the 3D points from the treehill dataset and initialize them as gaussian splats. In order to avoid a costly nearest neighbor search we initialize all scale variables as .01 (Note that with such a small variance we will need a strong concentration of splats in one spot to be visible. Larger variance makes the process quite slow.). Then all we have to do is call render_image with the image number we are trying to emulate and as you an see we get a sparse set of point clouds that resemble our image! (Check out our bonus section at the bottom for an equivalent CUDA kernel using pyTorch’s nifty tool that compiles CUDA code!)Actual image, CPU implementation, CUDA implementation. Image by author.While the backwards pass is not part of this tutorial, one note should be made that while we start with only these few points, we soon have hundreds of thousands of splats for most scenes. This is caused by the breaking up of large splats (as defined by larger variance on axes) into smaller splats and removing splats that have extremely low opacity. For instance, if we truly initialized the scale to the mean of the three closest nearest neighbors we would have a majority of the space covered. In order to get fine detail we would need to break these down into much smaller splats that are able to capture fine detail. They also need to populate areas with very few gaussians. They refer to these two scenarios as over reconstruction and under reconstruction and define both scenarios by large gradient values for various splats. They then split or clone the splats depending on size (see image below) and continue the optimization process.Although the backward pass is not covered in this tutorial, it’s important to note that we start with only a few points but soon have hundreds of thousands of splats in most scenes. This increase is due to the splitting of large splats (with larger variances on axes) into smaller ones and the removal of splats with very low opacity. For instance, if we initially set the scale to the mean of the three nearest neighbors, most of the space would be covered. To achieve fine detail, we need to break these large splats into much smaller ones. Additionally, areas with very few Gaussians need to be populated. These scenarios are referred to as over-reconstruction and under-reconstruction, characterized by large gradient values for various splats. Depending on their size, splats are split or cloned (see image below), and the optimization process continues.From the Author’s original paper on how gaussians are split or cloned in training. Source: https://arxiv.org/abs/2308.04079And that is an easy introduction to Gaussian Splatting! You should now have a good intuition on what exactly is going on in the forward pass of a gaussian scene render. While a bit daunting and not exactly neural networks, all it takes is a bit of linear algebra and we can render 3D geometry in 2D!Feel free to leave comments about confusing topics or if I got something wrong and you can always connect with me on LinkedIn or twitter!

Related articles

Recent articles