fix to update T when break

nerfstudio-project · Feb 10, 2025 · 1dc06b3 · 1dc06b3
1 parent 0880d2b
commit 1dc06b3
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 12 deletions.
diff --git a/gsplat/cuda/csrc/rasterize_to_pixels_2dgs_fwd.cu b/gsplat/cuda/csrc/rasterize_to_pixels_2dgs_fwd.cu
@@ -302,12 +302,6 @@ __global__ void rasterize_to_pixels_fwd_2dgs_kernel(
                 continue;
             }
 
-            const S next_T = T * (1.0f - alpha);
-            if (next_T <= 1e-4) { // this pixel is done: exclusive
-                done = true;
-                break;
-            }
-
             // run volumetric rendering..
             int32_t g = id_batch[t];
             const S vis = alpha * T;
@@ -344,6 +338,13 @@ __global__ void rasterize_to_pixels_fwd_2dgs_kernel(
 
             cur_idx = batch_start + t;
 
+            const S next_T = T * (1.0f - alpha);
+            if (next_T <= 1e-4) { // this pixel is done: exclusive
+                T = next_T;
+                done = true;
+                break;
+            }
+
             T = next_T;
         }
     }

diff --git a/gsplat/cuda/csrc/rasterize_to_pixels_fwd.cu b/gsplat/cuda/csrc/rasterize_to_pixels_fwd.cu
@@ -147,12 +147,6 @@ __global__ void rasterize_to_pixels_fwd_kernel(
                 continue;
             }
 
-            const S next_T = T * (1.0f - alpha);
-            if (next_T <= 1e-4) { // this pixel is done: exclusive
-                done = true;
-                break;
-            }
-
             int32_t g = id_batch[t];
             const S vis = alpha * T;
             const S *c_ptr = colors + g * COLOR_DIM;
@@ -162,6 +156,13 @@ __global__ void rasterize_to_pixels_fwd_kernel(
             }
             cur_idx = batch_start + t;
 
+            const S next_T = T * (1.0f - alpha);
+            if (next_T <= 1e-4) { // this pixel is done: exclusive
+                T = next_T;
+                done = true;
+                break;
+            }
+
             T = next_T;
         }
     }