add_body_envelope: add step-8 table detection to remove the CT table

scripts/utils.py

@@ -432,6 +432,7 @@ def add_body_envelope(
     hu_threshold: float = -800.0,
     closing_kernel: int = 3,
     bed_cleanup_kernel: int = 5,
+    table_frac_thresh: float = 0.05,
     device: str = "cuda:0",
 ):
     """
@@ -461,6 +462,10 @@ def add_body_envelope(
       6. **Force-include labels**: labeled voxels are forced inside body.
       7. **Fill**: every voxel inside the silhouette that the segmentation
          didn't already label is set to ``body_label``.
+      8. **Table detection** (safety net): if the air-density CT table leaked
+         into the body, it shows up as the largest connected component of
+         body voxels that are actually air (``CT < hu_threshold``); drop it
+         when that component is ``>= table_frac_thresh`` of the body.
 
     Args:
         seg_mask: ``(H, W, D)`` integer label volume (numpy ndarray or torch
@@ -477,6 +482,12 @@ def add_body_envelope(
             erode→LCC→dilate (step 4). Default 5 (slightly larger than
             ``closing_kernel`` so the body fully separates from the bed
             before the LCC selects it).
+        table_frac_thresh: step-8 table detector fires when the largest
+            air-in-body connected component (body voxels with
+            ``CT < hu_threshold``) is >= this fraction of the body, in which
+            case it is treated as the CT table and removed. Default 0.05 (a
+            table is empirically 16-28% of the body vs <0.3% clean, so any
+            value in ~0.02-0.10 separates them).
         device: torch device used for the morphology ops.
 
     Returns:
@@ -528,6 +539,28 @@ def add_body_envelope(
     body_np = body_t.detach().cpu().numpy() > 0.5
     out = seg_np.copy()
     out[body_np & (out == 0)] = body_label
+
+    # 8. Table detection. The find-air-invert steps above can still leak the air-density CT table into the body —
+    #    the air trapped between patient and table is a SEPARATE component from the exterior air, so it reads as
+    #    "not air" -> body. Detect it as the largest connected component of body voxels that are actually AIR
+    #    (``air_hu = CT < hu_threshold``); since the seg labels the lungs, no legitimate air region is anywhere
+    #    near table-sized (empirically a table is 16-28% of body vs <0.3% clean), so if the largest air-in-body
+    #    component is >= ``table_frac_thresh`` of the body, it's the table — drop it from the body.
+    air_hu = ct_np < hu_threshold  # air / low-density mask (same HU cut as the air step)
+    air_body = (out == body_label) & air_hu  # body voxels that are actually air
+    if air_body.any():
+        table = (
+            np.asarray(  # np.asarray for consistency with steps 1 & 4
+                get_largest_connected_component_mask(air_body.astype(np.float32), connectivity=None, num_components=1),
+                dtype=np.float32,
+            )
+            > 0.5
+        )
+        n_body = int((out == body_label).sum())
+        n_table = int(table.sum())
+        if n_body and n_table >= table_frac_thresh * n_body:
+            out[table] = 0  # remove the detected table from the body
+            print(f"[add_body_envelope] table detected ({100.0 * n_table / n_body:.1f}% of body) -> removed", flush=True)
     return out.astype(orig_dtype)