| |
| """ |
| Generate a Parquet dataset of TCGA patches described in sample_dataset_30.txt. |
| |
| Each line in the input text file is expected to follow the format emitted by |
| create_sample_dataset_txt.py: |
| |
| /path/to/slide.svs <x> <y> <level> |
| |
| Lines that include MPP metadata from the ablation sampler are also accepted: |
| |
| /path/to/slide.svs <x> <y> <level> <mpp_x> <mpp_y> |
| |
| The script opens every referenced slide once, extracts a 224x224 RGB patch at the |
| requested coordinates/level, and writes the samples into per-task Parquet files. |
| It supports PNG/JPEG/raw outputs, task chunking, resume markers, and optional |
| multi-processing. Results land in /data/TCGA_parquet_sample30/ unless overridden. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import atexit |
| import logging |
| import multiprocessing as mp |
| from collections import OrderedDict |
| from concurrent.futures import ProcessPoolExecutor, as_completed |
| from dataclasses import dataclass |
| from io import BytesIO |
| import os |
| import random |
| from pathlib import Path |
| from typing import Iterable, List, Sequence, Set, Tuple |
|
|
| import numpy as np |
| import pyarrow as pa |
| import pyarrow.parquet as pq |
| from openslide import OpenSlide |
|
|
|
|
| @dataclass(frozen=True) |
| class PatchSpec: |
| slide_path: str |
| x: int |
| y: int |
| level: int |
|
|
|
|
| @dataclass(frozen=True) |
| class TaskConfig: |
| tile_size: int |
| encoding: str |
| max_open_slides: int |
| output_dir: str |
| progress_dir: str |
| parquet_compression: str | None |
|
|
|
|
| CACHE_LIMIT_ENV = "PARQUET_MAX_OPEN_SLIDES" |
| try: |
| _env_cache_limit = int(os.environ.get(CACHE_LIMIT_ENV, "16")) |
| except (TypeError, ValueError): |
| _env_cache_limit = 16 |
| DEFAULT_MAX_OPEN_SLIDES = max(1, _env_cache_limit) |
| MAX_OPEN_SLIDES = DEFAULT_MAX_OPEN_SLIDES |
| SLIDE_CACHE: "OrderedDict[str, OpenSlide]" = OrderedDict() |
|
|
|
|
| def prune_slide_cache() -> None: |
| """Ensure the slide cache does not exceed MAX_OPEN_SLIDES.""" |
| if MAX_OPEN_SLIDES < 1: |
| return |
| while len(SLIDE_CACHE) > MAX_OPEN_SLIDES: |
| old_path, old_slide = SLIDE_CACHE.popitem(last=False) |
| try: |
| old_slide.close() |
| except Exception: |
| logging.exception("Failed to close slide %s during cache prune.", old_path) |
|
|
|
|
| def set_max_open_slides(limit: int) -> None: |
| """Update cache limit and prune if needed.""" |
| global MAX_OPEN_SLIDES |
| limit = max(1, int(limit)) |
| if limit == MAX_OPEN_SLIDES: |
| return |
| MAX_OPEN_SLIDES = limit |
| prune_slide_cache() |
|
|
|
|
| def ensure_cache_limit(limit: int) -> None: |
| """Idempotent helper invoked inside workers.""" |
| try: |
| set_max_open_slides(limit) |
| except Exception: |
| logging.exception("Unexpected failure updating slide cache limit.") |
|
|
|
|
| @atexit.register |
| def close_all_slides() -> None: |
| """Close any slides left in the cache when the process exits.""" |
| while SLIDE_CACHE: |
| path, slide = SLIDE_CACHE.popitem(last=False) |
| try: |
| slide.close() |
| except Exception: |
| logging.exception("Failed to close slide %s during interpreter shutdown.", path) |
|
|
|
|
| def get_slide(path: str) -> OpenSlide: |
| slide = SLIDE_CACHE.get(path) |
| if slide is None: |
| slide = OpenSlide(path) |
| SLIDE_CACHE[path] = slide |
| SLIDE_CACHE.move_to_end(path, last=True) |
| prune_slide_cache() |
| else: |
| SLIDE_CACHE.move_to_end(path, last=True) |
| return slide |
|
|
|
|
| def parse_specs(spec_file: Path) -> List[PatchSpec]: |
| specs: List[PatchSpec] = [] |
| with spec_file.open("r", encoding="utf-8") as fh: |
| for idx, raw_line in enumerate(fh, start=1): |
| line = raw_line.strip() |
| if not line or line.startswith("#"): |
| continue |
| parts = line.split() |
| if len(parts) == 4: |
| slide_path, x, y, level = parts |
| elif len(parts) == 6: |
| slide_path, x, y, level, *_mpp = parts |
| else: |
| raise ValueError(f"Invalid spec line {idx}: {line}") |
| specs.append(PatchSpec(slide_path, int(x), int(y), int(level))) |
| if not specs: |
| raise ValueError(f"No patch specifications found in {spec_file}") |
| logging.info("Parsed %d patch specifications from %s.", len(specs), spec_file) |
| return specs |
|
|
|
|
| def group_by_slide(specs: Sequence[PatchSpec]) -> List[Tuple[str, List[PatchSpec]]]: |
| """Group patch specs so each slide is opened only once.""" |
| grouped: "OrderedDict[str, List[PatchSpec]]" = OrderedDict() |
| for spec in specs: |
| grouped.setdefault(spec.slide_path, []).append(spec) |
| return [(slide, patches) for slide, patches in grouped.items()] |
|
|
|
|
| def build_tasks( |
| groups: Sequence[Tuple[str, Sequence[PatchSpec]]], |
| max_patches_per_task: int, |
| *, |
| start_index: int = 0, |
| ) -> List[Tuple[str, List[PatchSpec], str]]: |
| """Chunk grouped specs to improve parallel balance on large nodes.""" |
| tasks: List[Tuple[str, List[PatchSpec], str]] = [] |
| idx = start_index |
| for slide, patches in groups: |
| chunk = len(patches) if max_patches_per_task <= 0 else max_patches_per_task |
| for start in range(0, len(patches), chunk): |
| chunk_specs = list(patches[start : start + chunk]) |
| task_id = f"{idx:08d}" |
| tasks.append((slide, chunk_specs, task_id)) |
| idx += 1 |
| return tasks |
|
|
|
|
| def find_next_task_index(output_dir: Path) -> int: |
| """Return the next sequential task index based on existing Parquet files.""" |
| max_idx = -1 |
| for candidate in output_dir.glob("*.parquet"): |
| stem = candidate.stem |
| if stem.isdigit(): |
| try: |
| max_idx = max(max_idx, int(stem)) |
| except ValueError: |
| continue |
| return max_idx + 1 |
|
|
|
|
| def load_completed_task_ids(progress_dir: Path, output_dir: Path) -> Set[str]: |
| """Load completed task IDs from progress markers, pruning stale entries.""" |
| done_ids: Set[str] = set() |
| for marker in progress_dir.iterdir(): |
| if not marker.is_file(): |
| continue |
| if marker.name == "task_offset.txt": |
| continue |
| parquet_path = output_dir / f"{marker.stem}.parquet" |
| if parquet_path.exists(): |
| done_ids.add(marker.stem) |
| else: |
| marker.unlink() |
| return done_ids |
|
|
|
|
| def process_task_worker( |
| task: Tuple[str, Sequence[PatchSpec], str], |
| config: TaskConfig, |
| ) -> str: |
| """Process a single task and materialize it to a Parquet file.""" |
| slide_path, specs, task_id = task |
| ensure_cache_limit(config.max_open_slides) |
| slide = get_slide(slide_path) |
| buf = BytesIO() if config.encoding in ("png", "jpeg") else None |
|
|
| task_ids: List[str] = [] |
| slide_paths: List[str] = [] |
| xs: List[int] = [] |
| ys: List[int] = [] |
| levels: List[int] = [] |
| tile_sizes: List[int] = [] |
| level_downsamples: List[float] = [] |
| image_bytes: List[bytes] = [] |
| image_dtypes: List[str] = [] |
|
|
| for spec in specs: |
| if spec.level < 0 or spec.level >= slide.level_count: |
| raise ValueError(f"Level {spec.level} invalid for {slide_path}") |
| region = slide.read_region( |
| (spec.x, spec.y), spec.level, (config.tile_size, config.tile_size) |
| ).convert("RGB") |
|
|
| task_ids.append(task_id) |
| slide_paths.append(slide_path) |
| xs.append(spec.x) |
| ys.append(spec.y) |
| levels.append(spec.level) |
| tile_sizes.append(config.tile_size) |
| level_downsamples.append(float(slide.level_downsamples[spec.level])) |
|
|
| if config.encoding == "png": |
| buf.seek(0) |
| buf.truncate(0) |
| region.save(buf, format="PNG", optimize=True) |
| image_bytes.append(buf.getvalue()) |
| image_dtypes.append("uint8") |
| elif config.encoding == "jpeg": |
| buf.seek(0) |
| buf.truncate(0) |
| region.save(buf, format="JPEG", quality=95, optimize=True) |
| image_bytes.append(buf.getvalue()) |
| image_dtypes.append("uint8") |
| else: |
| arr = np.asarray(region, dtype=np.uint8) |
| image_bytes.append(arr.tobytes()) |
| image_dtypes.append(str(arr.dtype)) |
|
|
| table = pa.table( |
| { |
| "task_id": pa.array(task_ids, type=pa.string()), |
| "slide_path": pa.array(slide_paths, type=pa.string()), |
| "x": pa.array(xs, type=pa.int32()), |
| "y": pa.array(ys, type=pa.int32()), |
| "level": pa.array(levels, type=pa.int32()), |
| "tile_size": pa.array(tile_sizes, type=pa.int32()), |
| "level_downsample": pa.array(level_downsamples, type=pa.float32()), |
| "image_dtype": pa.array(image_dtypes, type=pa.string()), |
| "image_bytes": pa.array(image_bytes, type=pa.binary()), |
| } |
| ) |
|
|
| metadata = dict(table.schema.metadata or {}) |
| metadata.update( |
| { |
| b"image_encoding": config.encoding.encode("utf-8"), |
| b"tile_size": str(config.tile_size).encode("utf-8"), |
| } |
| ) |
| table = table.replace_schema_metadata(metadata) |
|
|
| output_path = Path(config.output_dir) / f"{task_id}.parquet" |
| pq.write_table(table, output_path, compression=config.parquet_compression) |
| progress_path = Path(config.progress_dir) / f"{task_id}.done" |
| progress_path.write_text("1\n") |
| return task_id |
|
|
|
|
| def parse_args(argv: Iterable[str] | None = None) -> argparse.Namespace: |
| parser = argparse.ArgumentParser( |
| description="Build a Parquet dataset of TCGA patches listed in a text file." |
| ) |
| parser.add_argument( |
| "--spec-file", |
| type=Path, |
| default=Path("sample_dataset_30.txt"), |
| help="Path to the text file produced by create_sample_dataset_txt.py.", |
| ) |
| parser.add_argument( |
| "--output-dir", |
| type=Path, |
| default=Path("/data/TCGA_parquet_sample30/"), |
| help="Directory where Parquet files will be written (default: /data/TCGA_parquet_sample30/).", |
| ) |
| parser.add_argument( |
| "--tile-size", |
| type=int, |
| default=224, |
| help="Square tile size (pixels) to read from each slide.", |
| ) |
| parser.add_argument( |
| "--num-workers", |
| type=int, |
| default=None, |
| help="Worker processes for parallel patch extraction (default: auto based on CPUs).", |
| ) |
| parser.add_argument( |
| "--start-method", |
| choices=("fork", "spawn", "forkserver"), |
| default=None, |
| help="Multiprocessing start method used by the worker pool (default: library default).", |
| ) |
| parser.add_argument( |
| "--mode", |
| choices=("overwrite", "append"), |
| default="overwrite", |
| help="Whether to overwrite or append to an existing Parquet dataset.", |
| ) |
| parser.add_argument( |
| "--keep-order", |
| action="store_true", |
| help="Keep samples in the same order as listed in the spec file.", |
| ) |
| parser.add_argument( |
| "--encoding", |
| choices=("png", "jpeg", "raw"), |
| default="png", |
| help="Patch serialization format (png: smaller, jpeg/raw: faster).", |
| ) |
| parser.add_argument( |
| "--task-chunk-size", |
| type=int, |
| default=42000, |
| help="Max patches per task to balance multi-process workloads (<=0 disables chunking). This is an arbitrarily high number which basically means that all patches for a slide will be processed in a single task. Can consider letting tasks span multiple slides in a future code revision if you want parquet files to be larger size.", |
| ) |
| parser.add_argument( |
| "--shuffle-tasks", |
| action="store_true", |
| help="Shuffle task order (ignored when --keep-order is set).", |
| ) |
| parser.add_argument( |
| "--resume", |
| action="store_true", |
| help="Skip tasks that completed in a previous run using progress markers.", |
| ) |
| parser.add_argument( |
| "--max-open-slides", |
| type=int, |
| default=DEFAULT_MAX_OPEN_SLIDES, |
| help=( |
| "Maximum slides cached per worker before least-recently-used eviction; " |
| ), |
| ) |
| parser.add_argument( |
| "--parquet-compression", |
| default="none", |
| choices=("none", "snappy", "gzip", "brotli", "zstd", "lz4"), |
| help="Compression codec for Parquet output (default: none).", |
| ) |
| parser.add_argument( |
| "--log-level", |
| default="INFO", |
| help="Python logging level (default: INFO).", |
| ) |
| return parser.parse_args(argv) |
|
|
|
|
| def main(argv: Iterable[str] | None = None) -> None: |
| args = parse_args(argv) |
| logging.basicConfig( |
| level=getattr(logging, args.log_level.upper(), logging.INFO), |
| format="%(asctime)s - %(levelname)s - %(message)s", |
| ) |
| ensure_cache_limit(args.max_open_slides) |
|
|
| if not args.spec_file.exists(): |
| raise SystemExit(f"Spec file not found: {args.spec_file}") |
|
|
| specs = parse_specs(args.spec_file) |
| if not specs: |
| raise SystemExit("No valid patch specs to process. Aborting.") |
|
|
| args.output_dir.mkdir(parents=True, exist_ok=True) |
|
|
| compression = None if args.parquet_compression == "none" else args.parquet_compression |
|
|
| existing_parquet = [ |
| path for path in args.output_dir.glob("*.parquet") if path.is_file() |
| ] |
| if args.mode == "overwrite" and not args.resume: |
| for path in existing_parquet: |
| try: |
| path.unlink() |
| except Exception: |
| logging.exception("Failed to remove %s during overwrite preparation.", path) |
| raise |
|
|
| progress_dir = args.output_dir / ".resume" |
| progress_dir.mkdir(parents=True, exist_ok=True) |
| if not args.resume: |
| for marker in progress_dir.iterdir(): |
| if marker.is_file(): |
| marker.unlink() |
|
|
| done_ids = load_completed_task_ids(progress_dir, args.output_dir) |
| offset_file = progress_dir / "task_offset.txt" |
|
|
| if args.resume: |
| if offset_file.exists(): |
| try: |
| start_index = int(offset_file.read_text().strip()) |
| except ValueError: |
| logging.warning("Ignoring malformed task offset metadata in %s.", offset_file) |
| start_index = 0 |
| elif done_ids: |
| try: |
| start_index = min(int(task_id) for task_id in done_ids) |
| except ValueError: |
| start_index = 0 |
| elif args.mode == "append": |
| start_index = find_next_task_index(args.output_dir) |
| else: |
| start_index = 0 |
| else: |
| start_index = find_next_task_index(args.output_dir) if args.mode == "append" else 0 |
|
|
| try: |
| offset_file.write_text(f"{start_index}\n") |
| except Exception: |
| logging.exception("Failed to persist task offset metadata to %s.", offset_file) |
| raise |
|
|
| grouped = group_by_slide(specs) |
| tasks = build_tasks( |
| grouped, |
| args.task_chunk_size, |
| start_index=start_index, |
| ) |
|
|
| if done_ids: |
| tasks = [task for task in tasks if task[2] not in done_ids] |
|
|
| if args.shuffle_tasks and not args.keep_order: |
| random.shuffle(tasks) |
|
|
| if not tasks: |
| logging.info("No tasks remaining after resume check.") |
| return |
|
|
| pending_patches = sum(len(task[1]) for task in tasks) |
| cpu_count = os.cpu_count() or 1 |
| num_workers = args.num_workers |
| if num_workers is None: |
| num_workers = max(1, min(cpu_count, 64)) |
|
|
| logging.info( |
| "Building Parquet dataset in %s (num_workers=%d, mode=%s, encoding=%s, compression=%s).", |
| args.output_dir, |
| num_workers, |
| args.mode, |
| args.encoding, |
| compression or "none", |
| ) |
| logging.info( |
| "Prepared %d task(s) across %d slide(s) [chunk_size=%d, pending_patches=%d].", |
| len(tasks), |
| len(grouped), |
| args.task_chunk_size, |
| pending_patches, |
| ) |
| if done_ids: |
| logging.info("Skipping %d previously completed task(s).", len(done_ids)) |
|
|
| config = TaskConfig( |
| tile_size=args.tile_size, |
| encoding=args.encoding, |
| max_open_slides=args.max_open_slides, |
| output_dir=str(args.output_dir), |
| progress_dir=str(progress_dir), |
| parquet_compression=compression, |
| ) |
|
|
| if num_workers <= 1: |
| for idx, task in enumerate(tasks, start=1): |
| task_id = process_task_worker(task, config) |
| logging.info( |
| "Completed task %s (%d/%d) for slide %s (%d patch(es)).", |
| task_id, |
| idx, |
| len(tasks), |
| task[0], |
| len(task[1]), |
| ) |
| else: |
| mp_context = mp.get_context(args.start_method) if args.start_method else None |
| with ProcessPoolExecutor(max_workers=num_workers, mp_context=mp_context) as executor: |
| futures = {executor.submit(process_task_worker, task, config): task for task in tasks} |
| for idx, future in enumerate(as_completed(futures), start=1): |
| slide_path, specs_for_task, task_id_hint = futures[future] |
| try: |
| completed_task_id = future.result() |
| except Exception: |
| logging.exception( |
| "Task %s failed for slide %s (%d patch(es)).", |
| task_id_hint, |
| slide_path, |
| len(specs_for_task), |
| ) |
| raise |
| logging.info( |
| "Completed task %s (%d/%d) for slide %s (%d patch(es)).", |
| completed_task_id, |
| idx, |
| len(tasks), |
| slide_path, |
| len(specs_for_task), |
| ) |
| logging.info("Dataset build complete.") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
|
|
| |
|
|
