Skip to content

Fix/tiling shm #7

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
kisside wants to merge 5 commits into
base: master
Choose a base branch
from
Draft

Fix/tiling shm #7

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
445203d
fix: add shm + conditional epithelium mask
kisside Mar 6, 2026
240c4bf
fix: lower memory, bigger repartition
kisside Mar 7, 2026
05fda03
fix: increase memory constrains
kisside Mar 9, 2026
8f8efd0
fix: change materialize order
kisside Mar 14, 2026
0230d97
fix: bump versions
kisside Apr 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 39 additions & 12 deletions preprocessing/tiling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,14 @@
from shapely.geometry import box
from shapely.geometry.polygon import Polygon


BATCH_SIZE = 256


def add_missing_epithelium_overlap(batch: DataBatch) -> DataBatch:
batch["epithelium_overlap"] = np.nan
return batch


def add_tile_overlap(
tiles: Dataset,
roi: Polygon,
Expand Down Expand Up @@ -78,7 +82,7 @@ def add_tile_overlap(
col("mpp_y"),
)
tiles = tiles.with_column(
overlap_struct_col, overlap_struct, num_cpus=1, memory=BATCH_SIZE * 8 * 512**2
overlap_struct_col, overlap_struct, num_cpus=0.2, memory=512 * 1024**2
)

def extract_value(batch: DataBatch) -> DataBatch:
Expand All @@ -93,8 +97,8 @@ def extract_value(batch: DataBatch) -> DataBatch:
tiles = tiles.map_batches(
extract_value,
batch_format="pandas",
num_cpus=1,
memory=BATCH_SIZE * 8 * 512**2,
num_cpus=0.1,
memory=256 * 1024**2,
)

return tiles.drop_columns([overlap_struct_col])
Expand Down Expand Up @@ -204,7 +208,12 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
)

# Create unique slide IDs and save slide metadata
slides = slides.map(row_hash, num_cpus=0.1, memory=128 * 1024**2)
slides = slides.map(row_hash, num_cpus=0.1, memory=256 * 1024**2)

# Materialize slides so the source pipeline (read + join + hash) runs once
# and is reused for both slides_df and the tiles pipeline.
slides = slides.materialize()
slides_df = slides.to_pandas()

# Expand slides into tile coordinates
tiles = slides.flat_map(
Expand All @@ -215,8 +224,8 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
epithelium_masks_path=epithelium_masks_path,
),
num_cpus=0.2,
memory=BATCH_SIZE * 1024**2,
).repartition(target_num_rows_per_block=BATCH_SIZE)
memory=256 * 1024**2,
).repartition(target_num_rows_per_block=4096)

# Compute masks and filter tissue tiles
# Tissue mask checks center 50%
Expand All @@ -231,16 +240,35 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
tiles, tissue_roi, "tissue_mask_path", "tissue_overlap", "0"
)
tiles = tiles.filter(expr=col("tissue_overlap") > 0.0)

# Materialize after tissue filter to:
# 1. Lock in the reduction (discard background tiles before computing remaining overlaps)
# 2. Break the long operator chain to prevent streaming backpressure starvation
tiles = tiles.materialize()

tiles = add_tile_overlap(tiles, full_roi, "blur_mask_path", "blur_overlap", "255")
tiles = add_tile_overlap(
tiles, full_roi, "folding_mask_path", "folding_overlap", "255"
)
tiles = add_tile_overlap(
tiles, full_roi, "residual_mask_path", "residual_overlap", "0"
)
tiles = add_tile_overlap(
tiles, full_roi, "epithelium_mask_path", "epithelium_overlap", "0"
)
if epithelium_masks_path:
print("[INFO] Computing overlap: epithelium_overlap from epithelium_mask_path")
tiles = add_tile_overlap(
tiles, full_roi, "epithelium_mask_path", "epithelium_overlap", "0"
)
else:
print(
"[INFO] Skipping epithelium overlap: dataset.mlflow_uris.epithelium_masks is not set"
)
Comment on lines +257 to +264

@gemini-code-assist gemini-code-assist Bot Apr 22, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

For better logging practices and consistency, it's recommended to use the logging module instead of print for informational messages. This allows for more flexible log level management and redirection.

You would need to add import logging at the top of the file. Then you can change the print statements to logging.info(...).


tiles = tiles.map_batches(
add_missing_epithelium_overlap,
batch_format="pandas",
num_cpus=0.1,
memory=256 * 1024**2,
)
Comment on lines +266 to +271

@gemini-code-assist gemini-code-assist Bot Apr 22, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Using map_batches just to add a column with a constant value is less efficient and idiomatic than using add_column. Ray Data's add_column is optimized for this use case and is more readable.

You can replace this map_batches call and the add_missing_epithelium_overlap function (lines 29-31) with a single line:

tiles = tiles.add_column("epithelium_overlap", lambda df: np.nan)

This change would make the code more concise and performant.

        tiles = tiles.add_column("epithelium_overlap", lambda df: np.nan)


# Drop unnecessary columns
tiles = tiles.drop_columns(
Expand All @@ -255,8 +283,7 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
]
)

# Convert Ray Datasets to pandas DataFrames
slides_df = slides.to_pandas()
# Convert tiles Ray Dataset to pandas DataFrame
tiles_df = tiles.to_pandas()

save_mlflow_dataset(
Expand Down
4 changes: 2 additions & 2 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,10 @@ dependencies = [
"hydra-core>=1.3.2",
"lightning>=2.5.5",
"omegaconf>=2.3.0",
"rationai-mlkit>=0.2.2",
"rationai-mlkit>=0.4.0",
"rationai-masks>=1.1.1",
"rationai-tiling>=1.1.1",
"ratiopath>=1.1.1",
"ratiopath>=1.3.1",
"rationai-sdk"
]

Expand Down
5 changes: 3 additions & 2 deletions scripts/preprocessing/run_tiling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,9 @@
job_name="mammaprint-tiling-...",
username=...,
image="cerit.io/rationai/base:2.0.6",
cpu=8,
memory="32Gi",
cpu=16,
memory="64Gi",
shm="32Gi",
gpu=None,
public=False,
script=[
Expand Down
Loading