diff --git a/configs/experiment/preprocessing/tile_masks_05mpp.yaml b/configs/experiment/preprocessing/tile_masks_05mpp.yaml new file mode 100644 index 0000000..8cfc5a9 --- /dev/null +++ b/configs/experiment/preprocessing/tile_masks_05mpp.yaml @@ -0,0 +1,16 @@ +# @package _global_ + +defaults: + - /data: dataset + - _self_ + +slides_artifact_path: train_split/slides.parquet +tiles_artifact_path: train_split/tiles.parquet +tile_percentage_cols: + - tile_coverage_Nerve + - tile_coverage_Blood + - tile_coverage_Connective-Tissue + - tile_coverage_Fat + - tile_coverage_Epithelium + - tile_coverage_Muscle + - tile_coverage_Other diff --git a/configs/preprocessing/tile_masks.yaml b/configs/preprocessing/tile_masks.yaml new file mode 100644 index 0000000..a287440 --- /dev/null +++ b/configs/preprocessing/tile_masks.yaml @@ -0,0 +1,12 @@ +# @package _global_ + +slides_artifact_path: ??? +tiles_artifact_path: ??? +tile_percentage_cols: ??? + +max_concurrent: 1 +mlflow_artifact_path: tile_masks + +metadata: + run_name: Tile Masks ${dataset.name} + description: Tile mask generation diff --git a/preprocessing/tile_masks.py b/preprocessing/tile_masks.py new file mode 100644 index 0000000..9738f3d --- /dev/null +++ b/preprocessing/tile_masks.py @@ -0,0 +1,107 @@ +from pathlib import Path +from tempfile import TemporaryDirectory +from typing import Any, cast + +import hydra +import mlflow +import pandas as pd +import pyvips +import torch +from omegaconf import DictConfig +from rationai.masks import tile_mask, write_big_tiff +from rationai.masks.mask_builders import ScalarMaskBuilder +from rationai.mlkit import autolog, with_cli_args +from rationai.mlkit.lightning.loggers import MLFlowLogger +from tqdm import tqdm + + +def process_slide( + slide: dict[str, Any], + output_dir: str, + tile_percentage_cols: list[str], + tiles_path: str, +) -> None: + slide_tiles = pd.read_parquet( + tiles_path, + columns=["x", "y", *tile_percentage_cols], + filters=[("slide_id", "==", slide["id"])], + ) + filename = f"{Path(slide['path']).stem}.tiff" + + for percentage_col in tile_percentage_cols: + builder = ScalarMaskBuilder( + save_dir=Path(output_dir, percentage_col), + filename=filename, + extent_x=slide["extent_x"], + extent_y=slide["extent_y"], + mpp_x=slide["mpp_x"], + mpp_y=slide["mpp_y"], + extent_tile=slide["tile_extent_x"], + stride=slide["stride_x"], + ) + builder.update( + data=torch.tensor(slide_tiles[percentage_col].values).unsqueeze(1), + xs=torch.tensor(slide_tiles["x"].values), + ys=torch.tensor(slide_tiles["y"].values), + ) + builder.save() + + mask = tile_mask( + slide_tiles, + tile_extent=(slide["tile_extent_x"], slide["tile_extent_y"]), + size=(slide["extent_x"], slide["extent_y"]), + ) + width, height = mask.size + mask_bytes = mask.tobytes() + del mask + + write_big_tiff( + image=pyvips.Image.new_from_memory( + data=mask_bytes, + width=width, + height=height, + bands=1, + format="uchar", + ), + path=Path(output_dir, "outlines", filename), + mpp_x=slide["mpp_x"], + mpp_y=slide["mpp_y"], + ) + + +@with_cli_args(["+preprocessing=tile_masks"]) +@hydra.main(config_path="../configs", config_name="preprocessing", version_base=None) +@autolog +def main(config: DictConfig, logger: MLFlowLogger) -> None: + tiling_run_id = config.dataset.mlflow_artifacts.tiling_run_id + tile_percentage_cols: list[str] = list(config.tile_percentage_cols) + + slides_path = mlflow.artifacts.download_artifacts( + run_id=tiling_run_id, + artifact_path=config.slides_artifact_path, + ) + tiles_path = mlflow.artifacts.download_artifacts( + run_id=tiling_run_id, + artifact_path=config.tiles_artifact_path, + ) + slides = pd.read_parquet(slides_path) + + items = cast("list[dict[str, Any]]", slides.to_dict(orient="records")) + + with TemporaryDirectory() as output_dir: + Path(output_dir, "outlines").mkdir() + for slide in tqdm(items): + process_slide( + slide, + output_dir=output_dir, + tile_percentage_cols=tile_percentage_cols, + tiles_path=tiles_path, + ) + + logger.log_artifacts( + local_dir=output_dir, artifact_path=config.mlflow_artifact_path + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/submit_tile_masks.py b/scripts/submit_tile_masks.py new file mode 100644 index 0000000..7fc8a94 --- /dev/null +++ b/scripts/submit_tile_masks.py @@ -0,0 +1,18 @@ +from kube_jobs import storage, submit_job + + +submit_job( + job_name="tissue-classification-tile-masks", + username=..., + cpu=8, + memory="64Gi", + gpu=None, + public=False, + script=[ + "git clone https://github.com/RationAI/tissue-classification.git workdir", + "cd workdir", + "uv sync", + "uv run python -m preprocessing.tile_masks +experiment=...", + ], + storage=[storage.secure.DATA], +)