diff --git a/scripts/utils.py b/scripts/utils.py index 4491dfa..a64e0f5 100644 --- a/scripts/utils.py +++ b/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)