DRDMsig/Data_Engineering

1"""2Detect and remove constant-value artifact planes at volume boundaries.3 4Interpolation during preprocessing can introduce planes filled with a single5non-zero constant value (e.g. 8.0 for CT) at the start or end of each spatial6axis.  This script:7 8  1. Scans all .nii.gz files under --image_dir (and optionally --label_dir).9  2. For each image, identifies boundary planes that are entirely one NON-ZERO10     value — zero-valued planes are skipped as they are legitimate background.11  3. Crops the artifact planes from image AND matching label (if present),12     preserving the spatial origin so the image stays in the same physical13     coordinate space.14  4. Overwrites in-place (use --dry-run to preview without writing).15 16Usage:17    # Dry-run: report artifacts without modifying files18    python clean_artifact_planes.py \19        --image_dir /path/to/MSD_processed/images \20        --label_dir /path/to/MSD_processed/labels \21        --dry-run22 23    # Actually clean:24    python clean_artifact_planes.py \25        --image_dir /path/to/MSD_processed/images \26        --label_dir /path/to/MSD_processed/labels27"""28import os29import glob30import argparse31import numpy as np32import SimpleITK as sitk33from tqdm import tqdm34 35 36def _get_plane(arr, axis, idx):37    """Extract a single plane from the array along the given axis."""38    slc = [slice(None)] * arr.ndim39    slc[axis] = idx40    return arr[tuple(slc)]41 42 43def find_artifact_slices(arr, axis, max_search=20):44    """Find contiguous constant-value boundary slices along `axis`.45 46    Returns (n_start, n_end): number of artifact slices to trim from the47    start and end of the given axis.48 49    A slice is considered an artifact if:50      - It has exactly 1 unique value, AND51      - That value is foreign to the adjacent interior plane (i.e. the value52        does not appear, or appears very rarely, in the neighbor).53    This avoids trimming legitimate background planes (e.g. -300 in CT air)54    that are naturally connected to interior regions with the same value.55    """56    n = arr.shape[axis]57 58    def _is_artifact(idx, interior_idx):59        plane = _get_plane(arr, axis, idx)60        unique = np.unique(plane)61        if len(unique) != 1:62            return False63        val = float(unique[0])64        # Check if this constant value appears in the adjacent interior plane65        neighbor = _get_plane(arr, axis, interior_idx)66        # If the value appears in >1% of the neighbor's voxels, it's likely67        # connected background, not an artifact68        match_ratio = np.mean(np.abs(neighbor - val) < 1e-6)69        if match_ratio > 0.01:70            return False71        return True72 73    # Find the first non-constant plane from each boundary to use as reference74    def _find_reference(start, stop, step):75        for idx in range(start, stop, step):76            plane = _get_plane(arr, axis, idx)77            if len(np.unique(plane)) > 1:78                return idx79        return start  # fallback80 81    ref_start = _find_reference(0, min(max_search + 5, n), 1)82    ref_end = _find_reference(n - 1, max(n - 1 - max_search - 5, -1), -1)83 84    n_start = 085    for i in range(min(max_search, n // 2)):86        if _is_artifact(i, ref_start):87            n_start = i + 188        else:89            break90 91    n_end = 092    for i in range(n - 1, max(n - 1 - max_search, n // 2), -1):93        if _is_artifact(i, ref_end):94            n_end = (n - 1 - i) + 195        else:96            break97 98    return n_start, n_end99 100 101def detect_artifacts(arr, max_search=20):102    """Detect artifact planes on all spatial axes.103 104    For 4D arrays (e.g. BRATS with shape [C, D, H, W]), only spatial axes105    (1, 2, 3) are checked; the channel axis (0) is skipped.106 107    Returns a dict: {axis: (n_start, n_end)} for axes that need trimming.108    """109    if arr.ndim == 3:110        spatial_axes = [0, 1, 2]111    elif arr.ndim == 4:112        spatial_axes = [1, 2, 3]113    else:114        spatial_axes = list(range(arr.ndim))115 116    crops = {}117    for axis in spatial_axes:118        n_start, n_end = find_artifact_slices(arr, axis, max_search=max_search)119        if n_start > 0 or n_end > 0:120            crops[axis] = (n_start, n_end)121    return crops122 123 124def build_crop_slices(ndim, crops):125    """Build a tuple of slices to crop the array according to `crops`."""126    slices = [slice(None)] * ndim127    for axis, (n_start, n_end) in crops.items():128        end = None if n_end == 0 else -n_end129        slices[axis] = slice(n_start, end)130    return tuple(slices)131 132 133def crop_sitk_image(sitk_img, crops):134    """Crop a SimpleITK image according to the detected artifact planes.135 136    Updates the origin so the cropped image occupies the correct physical space.137    """138    arr = sitk.GetArrayFromImage(sitk_img)139    crop_slices = build_crop_slices(arr.ndim, crops)140    cropped_arr = arr[crop_slices]141 142    cropped_img = sitk.GetImageFromArray(cropped_arr)143    cropped_img.SetSpacing(sitk_img.GetSpacing())144    cropped_img.SetDirection(sitk_img.GetDirection())145 146    # Adjust origin: SimpleITK arrays are in ZYX order, origin is in XYZ147    ndim_phys = sitk_img.GetDimension()  # physical dimensions (3 for 3D, 4 for 4D)148    origin = list(sitk_img.GetOrigin())149    spacing = list(sitk_img.GetSpacing())150    direction = np.array(sitk_img.GetDirection()).reshape(ndim_phys, ndim_phys)151 152    for axis, (n_start, _) in crops.items():153        if n_start > 0:154            # Map array axis to physical axis155            # SimpleITK: last array axis = first physical axis156            if arr.ndim == 3:157                phys_axis = 2 - axis158            elif arr.ndim == 4:159                phys_axis = 2 - (axis - 1)160            else:161                continue162            if phys_axis < ndim_phys:163                for i in range(min(3, ndim_phys)):164                    origin[i] += n_start * spacing[phys_axis] * direction[i, phys_axis]165 166    cropped_img.SetOrigin(origin)167 168    for key in sitk_img.GetMetaDataKeys():169        cropped_img.SetMetaData(key, sitk_img.GetMetaData(key))170 171    return cropped_img172 173 174def main():175    parser = argparse.ArgumentParser(description="Detect and remove constant-value artifact planes at volume boundaries.")176    parser.add_argument("--image_dir", type=str, required=True,177                        help="Directory containing .nii.gz image files.")178    parser.add_argument("--label_dir", type=str, default=None,179                        help="Directory containing matching .nii.gz label files (same filenames). "180                             "In recursive mode, labels are found at {subject_dir}/segmentation/{filename}.")181    parser.add_argument("--recursive", action="store_true",182                        help="Recursively search for .nii.gz files, excluding segmentation/ subdirs.")183    parser.add_argument("--max_search", type=int, default=20,184                        help="Max number of boundary slices to check per side (default: 20).")185    parser.add_argument("--dry-run", action="store_true",186                        help="Report artifacts without modifying any files.")187    args = parser.parse_args()188 189    if args.recursive:190        all_files = sorted(glob.glob(os.path.join(args.image_dir, "**", "*.nii.gz"), recursive=True))191        image_files = [f for f in all_files if '/segmentation/' not in f and '/label' not in f.lower()]192    else:193        image_files = sorted(glob.glob(os.path.join(args.image_dir, "*.nii.gz")))194    print(f"Found {len(image_files)} images in {args.image_dir}{' (recursive)' if args.recursive else ''}")195    if args.label_dir:196        print(f"Label dir: {args.label_dir}")197    if args.dry_run:198        print("*** DRY RUN — no files will be modified ***")199 200    total_artifacts = 0201    total_clean = 0202    total_slices_removed = 0203 204    for img_path in tqdm(image_files, desc="Scanning"):205        filename = os.path.basename(img_path)206        sitk_img = sitk.ReadImage(img_path)207        arr = sitk.GetArrayFromImage(sitk_img)208 209        crops = detect_artifacts(arr, max_search=args.max_search)210 211        if not crops:212            total_clean += 1213            continue214 215        total_artifacts += 1216        slices_removed = sum(s + e for s, e in crops.values())217        total_slices_removed += slices_removed218 219        detail = ", ".join(220            f"axis{ax}: -{s} start, -{e} end"221            for ax, (s, e) in sorted(crops.items())222        )223 224        # Report the artifact value225        for ax, (s, e) in crops.items():226            slc = [slice(None)] * arr.ndim227            if s > 0:228                slc[ax] = 0229            else:230                slc[ax] = arr.shape[ax] - 1231            val = arr[tuple(slc)].flat[0]232            break233        print(f"  {filename}: {arr.shape} -> trim {slices_removed} planes, val={val} ({detail})")234 235        if args.dry_run:236            continue237 238        # Crop and save image239        cropped_img = crop_sitk_image(sitk_img, crops)240        sitk.WriteImage(cropped_img, img_path)241 242        # Crop matching label if present243        if args.label_dir and not args.recursive:244            label_path = os.path.join(args.label_dir, filename)245            if os.path.isfile(label_path):246                sitk_lbl = sitk.ReadImage(label_path)247                cropped_lbl = crop_sitk_image(sitk_lbl, crops)248                sitk.WriteImage(cropped_lbl, label_path)249        elif args.recursive:250            # In recursive mode, look for label at {parent}/segmentation/{filename}251            parent_dir = os.path.dirname(img_path)252            label_path = os.path.join(parent_dir, 'segmentation', filename)253            if os.path.isfile(label_path):254                sitk_lbl = sitk.ReadImage(label_path)255                cropped_lbl = crop_sitk_image(sitk_lbl, crops)256                sitk.WriteImage(cropped_lbl, label_path)257 258    print(f"\nSummary:")259    print(f"  Total images:    {len(image_files)}")260    print(f"  With artifacts:  {total_artifacts}")261    print(f"  Clean:           {total_clean}")262    print(f"  Planes removed:  {total_slices_removed}")263    if args.dry_run:264        print("  (dry-run — nothing was modified)")265 266 267if __name__ == "__main__":268    main()269