diff --git a/scripts/infer_image_from_mask.py b/scripts/infer_image_from_mask.py index 68022b9..e9e5620 100644 --- a/scripts/infer_image_from_mask.py +++ b/scripts/infer_image_from_mask.py @@ -113,8 +113,7 @@ def ldm_conditional_sample_one_image_from_mask( combine_label = combine_label_or.to(device) if output_size[0] != combine_label.shape[2] or output_size[1] != combine_label.shape[3] or output_size[2] != combine_label.shape[4]: logging.info( - "output_size is not a desired value. Need to interpolate the mask to " - "match with output_size. The result image will be very low quality." + "output_size is not a desired value. Need to interpolate the mask to match with output_size. The result image will be very low quality." ) combine_label = torch.nn.functional.interpolate(combine_label, size=output_size, mode="nearest") diff --git a/scripts/utils.py b/scripts/utils.py index a64e0f5..6dfc5b2 100644 --- a/scripts/utils.py +++ b/scripts/utils.py @@ -28,7 +28,7 @@ from monai.transforms import Compose, EnsureTyped, Lambdad, LoadImaged, Orientationd from monai.transforms.utils_morphological_ops import dilate, erode from monai.utils import TransformBackends, convert_data_type, convert_to_dst_type, get_equivalent_dtype -from scipy import stats +from scipy import ndimage, stats from torch import Tensor @@ -433,6 +433,7 @@ def add_body_envelope( closing_kernel: int = 3, bed_cleanup_kernel: int = 5, table_frac_thresh: float = 0.05, + seg_has_lung: bool = True, device: str = "cuda:0", ): """ @@ -444,11 +445,7 @@ def add_body_envelope( but never the body envelope, so users must add it before running ``ldm_conditional_sample_one_image_from_mask``. This helper does that. - Algorithm follows ``find_body_maskv2`` from pengfeig's ``3d_ldm_monai`` - (find-air-then-invert with a two-stage bed/table cleanup), which is - more robust than naively largest-CC'ing the body voxels directly — - the patient bed often touches the body, so a simple largest-CC keeps - it. Steps: + Steps (find-air-then-invert with a two-stage bed/table cleanup): 1. **Air mask**: largest connected component of voxels with ``CT < hu_threshold`` (default -800 HU). Closed via dilate→erode @@ -463,9 +460,12 @@ def add_body_envelope( 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. + into the body, it shows up as connected components of body voxels + that are actually air (``CT < hu_threshold``); drop EVERY such + component that is ``>= table_frac_thresh`` of the body. A table can + split into multiple pieces (side rails / pads / broken by the patient + silhouette), so all table-sized components are removed, not just the + largest. Args: seg_mask: ``(H, W, D)`` integer label volume (numpy ndarray or torch @@ -482,12 +482,16 @@ 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). + table_frac_thresh: step-8 removes EVERY air-in-body connected component + (body voxels with ``CT < hu_threshold``) that is >= this fraction + of the body, treating each as CT table. Default 0.05 (a table piece + is empirically 8-28% of the body vs <0.3% clean, so any value in + ~0.02-0.10 separates them). Multiple table-sized components are all + removed (a table often splits into rails/pads). + seg_has_lung: whether ``seg_mask`` labels the lungs. True (default) runs + step 8 (table removal). Set False when the seg has no lungs — step 8 + is then skipped, since lung air would otherwise be indistinguishable + from the table and wrongly removed. device: torch device used for the morphology ops. Returns: @@ -542,25 +546,24 @@ def add_body_envelope( # 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 - ) + # "not air" -> body. Detect it as connected components of body voxels that are actually AIR + # (``air_hu = CT < hu_threshold``); since the seg labels the lungs, no legitimate air region is table-sized, + # so EVERY air-in-body component >= ``table_frac_thresh`` of the body is table. There can be MORE THAN ONE + # (a split table, separate side rails, or a table broken by the patient silhouette), so remove ALL of them. + if seg_has_lung: # only safe when lungs are labeled (lung air would otherwise look like the table) + 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 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) + if air_body.any() and n_body: + lbl, ncc = ndimage.label(air_body) # all air-in-body components + if ncc: + sizes = np.bincount(lbl.ravel()) + sizes[0] = 0 # drop background + table_ids = np.nonzero(sizes >= table_frac_thresh * n_body)[0] # every table-sized component + if table_ids.size: + out[np.isin(lbl, table_ids)] = 0 # remove them all from the body + fracs = [round(float(100.0 * sizes[i] / n_body), 1) for i in table_ids] + print(f"[add_body_envelope] table detected ({table_ids.size} component(s): {fracs}% of body) -> removed", flush=True) return out.astype(orig_dtype)