build_1ch_cache.py

"""
Build 1-channel CT tensor cache from DICOM files.

HU windowing: brain/blood window, clip [-300, 180], normalize to [0, 1].
Output: .npz files with key 'image_norm' (float16, 512×512).

Resume logic: skips files whose image_id already exists in output dir.
Bad files: logged to BadDicomFiles.log in output dir.
"""

import os
import sys
import time
import argparse
import numpy as np
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed
from datetime import datetime

# Directory containing dicom_reader_1ch.py (this file's directory)
NEWICH_DIR = os.path.dirname(os.path.abspath(__file__))

OUTPUT_DIR_DEFAULT = '/home/johnb/cache_1ch'

HU_LOW = -300.0
HU_HIGH = 180.0
HU_RANGE = HU_HIGH - HU_LOW   # 480.0

TARGET_SIZE = 512


def _process_one(args_tuple):
    """Worker function: read DICOM, preprocess, return (image_id, result_array, error_str)."""
    file_path, = args_tuple

    # Import locally for worker safety; use robust 1ch reader
    import sys
    sys.path.insert(0, NEWICH_DIR)
    from dicom_reader_1ch import read_dicom_hu, extract_image_id

    image_id = extract_image_id(file_path)

    try:
        hu, image_id, is_valid = read_dicom_hu(file_path)
    except Exception as e:
        return image_id, None, f"read_dicom_hu exception: {e}"

    if not is_valid or hu is None:
        return image_id, None, "invalid (read_dicom_hu is_valid=False)"

    # Resize to 512×512 if needed
    h, w = hu.shape
    if h != TARGET_SIZE or w != TARGET_SIZE:
        try:
            from scipy.ndimage import zoom
            zh, zw = TARGET_SIZE / h, TARGET_SIZE / w
            hu = zoom(hu, (zh, zw), order=1)
        except Exception as e:
            return image_id, None, f"resize error: {e}"

    # Clip to brain/blood window
    hu = np.clip(hu, HU_LOW, HU_HIGH)

    # Normalize to [0, 1]
    image_norm = ((hu - HU_LOW) / HU_RANGE).astype(np.float16)

    return image_id, image_norm, None


def scan_dcm_files(dcm_dir: str):
    """Recursively yield all .dcm file paths under dcm_dir."""
    for dirpath, _, filenames in os.walk(dcm_dir):
        for fname in filenames:
            if fname.lower().endswith('.dcm'):
                yield os.path.join(dirpath, fname)


def build_cache(args):
    os.makedirs(args.output_dir, exist_ok=True)
    log_path = os.path.join(args.output_dir, 'BadDicomFiles.log')

    # Build resume set
    existing_ids = set()
    with os.scandir(args.output_dir) as it:
        for entry in it:
            if entry.name.endswith('.npz') and entry.is_file():
                existing_ids.add(entry.name[:-4])   # strip .npz
    print(f"Resume set: {len(existing_ids)} already cached")

    # Collect all DICOM paths, filter against resume set
    print(f"Scanning DICOM directory: {args.dcm_dir}")
    all_files = []
    for fpath in scan_dcm_files(args.dcm_dir):
        # Quick ID check using stem (avoids importing dicom_reader in main process)
        stem = Path(fpath).stem
        # Crude image_id prediction: if it starts with ID_ use stem, otherwise skip
        # We'll do the real extraction in the worker; only exclude obvious hits
        if stem in existing_ids:
            continue
        all_files.append(fpath)

    print(f"Files to process (after quick resume filter): {len(all_files)}")

    written = 0
    skipped = 0
    bad = 0
    start_time = time.perf_counter()
    last_milestone = 0

    batch_size = args.batch_size
    total = len(all_files)

    with open(log_path, 'a') as log_fh:
        with ProcessPoolExecutor(max_workers=args.workers) as executor:
            # Submit in batches to limit memory of pending futures
            offset = 0
            while offset < total:
                batch = all_files[offset:offset + batch_size]
                offset += batch_size

                futures = {executor.submit(_process_one, (fp,)): fp
                           for fp in batch}

                for fut in as_completed(futures):
                    try:
                        image_id, image_norm, error = fut.result()
                    except Exception as exc:
                        fp = futures[fut]
                        ts = datetime.now().isoformat(timespec='seconds')
                        log_fh.write(f"{ts}\t{fp}\texecutor exception: {exc}\n")
                        bad += 1
                        continue

                    if error is not None:
                        fp = futures[fut]
                        ts = datetime.now().isoformat(timespec='seconds')
                        log_fh.write(f"{ts}\t{fp}\t{error}\n")
                        bad += 1
                        continue

                    # Check precise resume (worker extracted real image_id)
                    if image_id in existing_ids:
                        skipped += 1
                        continue

                    # Save
                    out_path = os.path.join(args.output_dir, f"{image_id}.npz")
                    try:
                        np.savez_compressed(out_path, image_norm=image_norm)
                        existing_ids.add(image_id)
                        written += 1
                    except Exception as exc:
                        fp = futures[fut]
                        ts = datetime.now().isoformat(timespec='seconds')
                        log_fh.write(f"{ts}\t{fp}\tsave error: {exc}\n")
                        bad += 1
                        continue

                    # Progress milestone every 10,000 written
                    milestone = written // 10000
                    if milestone > last_milestone:
                        last_milestone = milestone
                        elapsed = time.perf_counter() - start_time
                        print(f"[{written}] Written {written} tensors "
                              f"({skipped} skipped, {bad} bad) — {elapsed:.1f}s")
                        sys.stdout.flush()

    elapsed = time.perf_counter() - start_time
    print(f"\nDone. Written: {written}, Skipped: {skipped}, Bad: {bad}, "
          f"Elapsed: {elapsed:.1f}s")
    print(f"Bad file log: {log_path}")


def main():
    parser = argparse.ArgumentParser(
        description='Build 1-channel CT .npz cache from DICOM files')
    parser.add_argument('--dcm-dir', type=str, required=True,
                        help='DICOM root directory (recursive scan)')
    parser.add_argument('--output-dir', type=str, default=OUTPUT_DIR_DEFAULT,
                        help='Output cache directory (default: /home/johnb/cache_1ch)')
    parser.add_argument('--workers', type=int, default=8,
                        help='ProcessPoolExecutor workers (default: 8)')
    parser.add_argument('--batch-size', type=int, default=1000,
                        help='Files submitted to executor at once (default: 1000)')
    args = parser.parse_args()

    build_cache(args)


if __name__ == '__main__':
    main()