diff --git a/code/preprocessing/01_scanDicom.py b/code/preprocessing/01_scanDicom.py index cbe5b75..8378842 100755 --- a/code/preprocessing/01_scanDicom.py +++ b/code/preprocessing/01_scanDicom.py @@ -208,8 +208,9 @@ def _has_dcm_magic(path: str) -> bool: # Pipeline functions # --------------------------------------------------------------------------- -def _extractDicom_impl(f: str, logger: logging.Logger, slice_counts: Dict[str, int] = None) -> Optional[Dict[str, Any]]: +def _extractDicom_impl(f: str, slice_counts: Dict[str, int] = None) -> Optional[Dict[str, Any]]: """Extract DICOM information from a specific file path.""" + logger = logging.getLogger('01_scanDicom') try: logger.debug(f'Extracting information for file: {f}') directory = os.path.dirname(f) @@ -391,29 +392,53 @@ def _scan_subdir_worker(subdir: str, sample_pct: float, sample_seed: Optional[in """Worker for multiprocessing directory scanning. Calls `_find_dicom_worker` on its assigned single top-level subdirectory. - Creates its own logger inside the child process to avoid pickle-lock deadlock. Returns list of (dicom_files, slice_counts). """ - wlogger = logging.getLogger(__name__ + '.worker') - return [_find_dicom_worker(subdir, sample_pct, sample_seed, wlogger)] + return [_find_dicom_worker(subdir, sample_pct, sample_seed)] -def _scan_subdir(topdir: str): - """Return a list of subdirectories containing .dcm files.""" - dirs_with_dcm = [] - for root, _, files in os.walk(topdir, followlinks=False): - if any(f.lower().endswith('.dcm') for f in files): - dirs_with_dcm.append(root) - return dirs_with_dcm +def _scan_subdir(topdir: str, min_targets: int = 16): + """Return a list of disjoint subdirectories that cover the entire tree. + We gather directories using BFS until we have enough targets. If a directory contains + files directly, we stop expanding it to keep subtrees disjoint for the os.walk workers.""" + dirs_to_scan = [] + queue = [topdir] + + while queue and (len(dirs_to_scan) + len(queue)) < min_targets: + curr = queue.pop(0) + try: + with os.scandir(curr) as it: + subdirs = [] + has_files = False + for entry in it: + if entry.is_dir(follow_symlinks=False): + subdirs.append(entry.path) + elif entry.is_file(): + if entry.name.lower().endswith('.dcm'): + has_files = True + + if has_files: + dirs_to_scan.append(curr) + else: + queue.extend(subdirs) + except Exception: + pass + + dirs_to_scan.extend(queue) + + if not dirs_to_scan: + dirs_to_scan = [topdir] + + return dirs_to_scan -def _find_dicom_worker(directory: str, sample_pct: float, sample_seed: Optional[int], - logger: logging.Logger) -> tuple: +def _find_dicom_worker(directory: str, sample_pct: float, sample_seed: Optional[int]) -> tuple: """Worker for findDicom — called per directory, accepts only plain args. Returns: (dicom_files, slice_counts) """ + logger = logging.getLogger('01_scanDicom') dicom_files = [] slice_counts = {} @@ -524,20 +549,20 @@ def main(cfg: ScanConfig, logger: logging.Logger, out_name: str = 'Data_table.cs if combined_result is None: if cfg.parallel: - # Parallel scan: walk the tree once to find subdirectories with .dcm files, - # then dispatch multiprocessing workers across them. - dirs_with_dcm = _scan_subdir(scan_dir) - logger.info(f'Found {len(dirs_with_dcm)} directories with DICOM files to scan') + # Parallel scan: we parallelize the walk by getting immediate subdirectories + # and dispatching multiprocessing workers to walk those subtrees independently. + target_dirs = _scan_subdir(scan_dir, min_targets=cfg.n_cpus * 4) + logger.info(f'Found {len(target_dirs)} branch directories to scan in parallel') - if len(dirs_with_dcm) > 1: + if len(target_dirs) > 1: worker_results = run_function( - logger, _scan_subdir_worker, dirs_with_dcm, + logger, _scan_subdir_worker, target_dirs, Parallel=True, P_type='hybrid', N_CPUS=cfg.n_cpus, sample_pct=cfg.sample_pct, sample_seed=cfg.sample_seed, ) else: worker_results = [run_function( - logger, _scan_subdir_worker, dirs_with_dcm, + logger, _scan_subdir_worker, target_dirs, Parallel=False, P_type='thread', N_CPUS=1, sample_pct=cfg.sample_pct, sample_seed=cfg.sample_seed, )] @@ -580,7 +605,7 @@ def main(cfg: ScanConfig, logger: logging.Logger, out_name: str = 'Data_table.cs worker_results = run_function( logger, _find_dicom_worker, dicom_dirs, Parallel=cfg.parallel, P_type='hybrid', N_CPUS=cfg.n_cpus, - sample_pct=cfg.sample_pct, sample_seed=cfg.sample_seed, logger=logger, + sample_pct=cfg.sample_pct, sample_seed=cfg.sample_seed, ) dicom_files = [f for files, _ in worker_results for f in files] slice_counts = {} @@ -605,7 +630,7 @@ def main(cfg: ScanConfig, logger: logging.Logger, out_name: str = 'Data_table.cs info_list = None if info_list is None: - extract_partial = partial(_extractDicom_impl, logger=logger, slice_counts=slice_counts) + extract_partial = partial(_extractDicom_impl, slice_counts=slice_counts) info_list = run_function( logger, extract_partial, dicom_files, Parallel=cfg.parallel, P_type='hybrid', N_CPUS=cfg.n_cpus, diff --git a/test/test_scanDicom_full.py b/test/test_scanDicom_full.py index 7bb0471..03bed06 100644 --- a/test/test_scanDicom_full.py +++ b/test/test_scanDicom_full.py @@ -232,7 +232,7 @@ def test_A4_missing_series_number_no_crash(tmp_path): d.mkdir() make_realistic_mr_dcm(str(d / "ns.dcm"), modality='MR', series_number=1) logger = _scan_logger() - found_files, _ = scan._find_dicom_worker(str(d), sample_pct=0.0, sample_seed=None, logger=logger) + found_files, _ = scan._find_dicom_worker(str(d), sample_pct=0.0, sample_seed=None) assert isinstance(found_files, list) @@ -243,7 +243,7 @@ def test_A5_duplicate_series_returns_one(tmp_path): for i in range(5): make_minimal_dcm(str(root / f"dup_{i}.dcm"), modality='MR', series_number=42) logger = _scan_logger() - found_files, _ = scan._find_dicom_worker(str(root), sample_pct=0.0, sample_seed=None, logger=logger) + found_files, _ = scan._find_dicom_worker(str(root), sample_pct=0.0, sample_seed=None) assert len(found_files) == 1 @@ -256,7 +256,7 @@ def test_A6_corrupt_files(tmp_path): (d / "bad2.dcm").write_bytes(b'\xff' * 512) (d / "bad3.dcm").write_bytes(b'\0' * 100) logger = _scan_logger() - found_files, _ = scan._find_dicom_worker(str(d), sample_pct=0.0, sample_seed=None, logger=logger) + found_files, _ = scan._find_dicom_worker(str(d), sample_pct=0.0, sample_seed=None) assert len(found_files) == 1 assert "good.dcm" in found_files[0] @@ -278,8 +278,8 @@ def test_A8_sampling_deterministic(tmp_path): for i in range(20): make_minimal_dcm(str(root / f"f_{i:02d}.dcm"), modality='MR', series_number=(i % 5) + 1) logger = _scan_logger() - first = scan._find_dicom_worker(str(root), sample_pct=15.0, sample_seed=99, logger=logger) - second = scan._find_dicom_worker(str(root), sample_pct=15.0, sample_seed=99, logger=logger) + first = scan._find_dicom_worker(str(root), sample_pct=15.0, sample_seed=99) + second = scan._find_dicom_worker(str(root), sample_pct=15.0, sample_seed=99) assert first == second @@ -320,7 +320,7 @@ def test_B1_extractDicom_has_all_keys(tmp_path): f = tmp_path / "extract_test.dcm" make_realistic_mr_dcm(str(f), repetition_time=500.0) logger = _scan_logger() - result = scan._extractDicom_impl(str(f), logger) + result = scan._extractDicom_impl(str(f)) assert result is not None assert isinstance(result, dict) assert EXPECTED_KEYS.issubset(result.keys()), f"Missing keys: {EXPECTED_KEYS - result.keys()}" @@ -331,21 +331,21 @@ def test_B2_T1_vs_T2_modality(tmp_path): logger = _scan_logger() t1_path = tmp_path / "t1.dcm" make_realistic_mr_dcm(str(t1_path), repetition_time=779.0) - t1_result = scan._extractDicom_impl(str(t1_path), logger) + t1_result = scan._extractDicom_impl(str(t1_path)) assert t1_result['Modality'] == 'T1', f"Expected T1, got {t1_result['Modality']}" t2_path = tmp_path / "t2.dcm" make_realistic_mr_dcm(str(t2_path), repetition_time=780.0) - t2_result = scan._extractDicom_impl(str(t2_path), logger) + t2_result = scan._extractDicom_impl(str(t2_path)) assert t2_result['Modality'] == 'T2', f"Expected T2, got {t2_result['Modality']}" t1_edge = tmp_path / "t1_edge.dcm" make_realistic_mr_dcm(str(t1_edge), repetition_time=779.999) - assert scan._extractDicom_impl(str(t1_edge), logger)['Modality'] == 'T1' + assert scan._extractDicom_impl(str(t1_edge))['Modality'] == 'T1' t2_edge = tmp_path / "t2_edge.dcm" make_realistic_mr_dcm(str(t2_edge), repetition_time=780.001) - assert scan._extractDicom_impl(str(t2_edge), logger)['Modality'] == 'T2' + assert scan._extractDicom_impl(str(t2_edge))['Modality'] == 'T2' # B3 — Unknown fields for missing tags @@ -354,7 +354,7 @@ def test_B3_unknown_fields_missing_tags(tmp_path): d.mkdir() make_minimal_dcm(str(d / "sparse.dcm"), modality='MR', series_number=1) logger = _scan_logger() - result = scan._extractDicom_impl(str(d / "sparse.dcm"), logger) + result = scan._extractDicom_impl(str(d / "sparse.dcm")) assert result is not None for key in ['Accession', 'DOB', 'Lat']: assert result[key] == 'Unknown', f"{key} should be 'Unknown' but is '{result[key]}'" diff --git a/test/test_scanDicom_unit.py b/test/test_scanDicom_unit.py index 944e6e6..de43660 100644 --- a/test/test_scanDicom_unit.py +++ b/test/test_scanDicom_unit.py @@ -95,7 +95,7 @@ def test_findDicom_series(tmp_path): make_minimal_dcm(str(root / "b.dcm"), modality='MR', series_number=2) make_minimal_dcm(str(root / "c.dcm"), modality='CT', series_number=3) logger = _make_logger() - found_files, _ = scan._find_dicom_worker(str(root), sample_pct=0.0, sample_seed=None, logger=logger) + found_files, _ = scan._find_dicom_worker(str(root), sample_pct=0.0, sample_seed=None) assert any("a.dcm" in f or "b.dcm" in f for f in found_files) @@ -103,7 +103,7 @@ def test_extractDicom_basic(tmp_path): f = tmp_path / "x.dcm" make_minimal_dcm(str(f), modality='MR', series_number=5, patient_id='P1') logger = _make_logger() - out = scan._extractDicom_impl(str(f), logger) + out = scan._extractDicom_impl(str(f)) assert isinstance(out, dict) assert isinstance(out['Modality'], str) @@ -125,7 +125,7 @@ def test_findDicom_handles_unreadable_and_returns_mr_only(tmp_path): make_minimal_dcm(str(root / "mri.dcm"), modality='MR', series_number=10) (root / "garbage.dcm").write_text("corrupt") logger = _make_logger() - found_files, _ = scan._find_dicom_worker(str(root), sample_pct=0.0, sample_seed=None, logger=logger) + found_files, _ = scan._find_dicom_worker(str(root), sample_pct=0.0, sample_seed=None) assert any("mri.dcm" in f for f in found_files) @@ -137,6 +137,6 @@ def test_findDicom_sampling_is_deterministic_with_seed(tmp_path): make_minimal_dcm(str(root / f"img_{i}.dcm"), modality='MR', series_number=series) logger = _make_logger() - first = scan._find_dicom_worker(str(root), sample_pct=20.0, sample_seed=123, logger=logger) - second = scan._find_dicom_worker(str(root), sample_pct=20.0, sample_seed=123, logger=logger) + first = scan._find_dicom_worker(str(root), sample_pct=20.0, sample_seed=123) + second = scan._find_dicom_worker(str(root), sample_pct=20.0, sample_seed=123) assert first == second \ No newline at end of file