From 80e2b8f3284ee17a7d1715740fd4c372d977d5fd Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Thu, 14 May 2026 13:05:43 +0200
Subject: [PATCH 01/23] feat: add provgigapath

---
 docs/available-models.md        | 27 ++++++++++++++++++++++++++-
 docs/guides/deployment-guide.md |  2 ++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/docs/available-models.md b/docs/available-models.md
index ca7a0e1..dc593ff 100644
--- a/docs/available-models.md
+++ b/docs/available-models.md
@@ -11,6 +11,7 @@ All endpoints receive and return data over HTTP using `POST` requests. To minimi
 | **Prostate Classifier 1** | `/prostate-classifier-1` | Binary Classification         |
 | **Episeg 1**              | `/episeg-1`              | Semantic Segmentation         |
 | **Virchow2**              | `/virchow2`              | Foundation Model / Embeddings |
+| **Prov-GigaPath**         | `/prov-gigapath`         | Foundation Model / Embeddings |
 | **Heatmap Builder**       | `/heatmap-builder`       | Pipeline / Custom Builder     |
 
 ---
@@ -80,7 +81,31 @@ with Client() as client:
   print(emb.shape)
 ```
 
-### 4. Heatmap Builder (`/heatmap-builder`)
+### 4. Prov-GigaPath (`/prov-gigapath`)
+
+A foundation model for pathology tile embeddings (Prov-GigaPath).
+
+- **Input**: LZ4-compressed raw bytes of a tissue tile image (`uint8`, shape `(tile_size, tile_size, 3)`).
+- **Output**: Output tensor matching the user's requested precision.
+- **Headers**:
+  - `x-output-dtype` (optional, default: `float32`): Sets the return precision (for example `float16`).
+- **SDK example**:
+
+```python
+from rationai import Client
+import numpy as np
+
+with Client() as client:
+  emb = client.models.embed_image(
+    model="prov-gigapath",
+    image=image,
+    output_dtype=np.float16,
+    timeout=30.0,
+  )
+  print(emb.shape)
+```
+
+### 5. Heatmap Builder (`/heatmap-builder`)
 
 A processing pipeline element for aggregating inferences into spatial heatmaps.
 
diff --git a/docs/guides/deployment-guide.md b/docs/guides/deployment-guide.md
index 56dc372..6dad494 100644
--- a/docs/guides/deployment-guide.md
+++ b/docs/guides/deployment-guide.md
@@ -49,6 +49,8 @@ Create a file in `helm/rayservice/applications/` (for example `my-model.yaml`) w
         max_replicas: 4
 ```
 
+- Add the new application file name to `helm/rayservice/values.yaml` under `applications`.
+
 Notes:
 
 - Use a dedicated branch in `working_dir` during development.

From c4374d7a3400abb39186c242c1fafb7ab9a40418 Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Thu, 14 May 2026 13:35:38 +0200
Subject: [PATCH 02/23] feat: new heatmap

Co-authored-by: Copilot <copilot@github.com>
---
 builders/heatmap_builder.py  | 66 ++++++++++++++++++++++++++-----
 misc/tile_heatmap_builder.py | 75 ------------------------------------
 2 files changed, 56 insertions(+), 85 deletions(-)
 delete mode 100644 misc/tile_heatmap_builder.py

diff --git a/builders/heatmap_builder.py b/builders/heatmap_builder.py
index 5b40636..d44a5ab 100644
--- a/builders/heatmap_builder.py
+++ b/builders/heatmap_builder.py
@@ -1,5 +1,6 @@
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
 from typing import Any, TypedDict
 
 from fastapi import FastAPI
@@ -35,11 +36,13 @@ async def root(
         output_bigtiff_tile_height: int,
         output_bigtiff_tile_width: int,
     ) -> str:
+        import numpy as np
+        import pyvips
+        from ratiopath.masks.mask_builders.mask_builder import MaskBuilder
         from ratiopath.openslide import OpenSlide
         from ratiopath.tiling import grid_tiles
 
         from misc.fetch_tissue_tile import fetch_tissue_tile
-        from misc.tile_heatmap_builder import TileHeatmapBuilder
 
         model = serve.get_app_handle(model_id)
         model_config = await model.get_config.remote()
@@ -64,11 +67,11 @@ async def root(
             scale_x = tissue_extent_x / extent_x
             scale_y = tissue_extent_y / extent_y
 
-            mask_builder = TileHeatmapBuilder(
-                extent_x=extent_x, extent_y=extent_y, mpp_x=mpp_x, mpp_y=mpp_y
-            )
+            mask_builder: MaskBuilder | None = None
+            mask_builder_lock = asyncio.Lock()
 
             async def process_tile(x: int, y: int) -> None:
+                nonlocal mask_builder
                 tile = await loop.run_in_executor(
                     executor,
                     fetch_tissue_tile,
@@ -82,9 +85,40 @@ async def process_tile(x: int, y: int) -> None:
                     tissue_level,
                     model_config["tile_size"],
                 )
-                if tile is not None:
-                    prediction = await model.predict.remote(tile)
-                    mask_builder.update(prediction, x, y)
+                if tile is None:
+                    return
+
+                prediction = await model.predict.remote(tile)
+                arr = np.asarray(prediction)
+
+                # Normalize to (B, C, *spatial) shape
+                match arr.ndim:
+                    case 0:
+                        batch = arr.reshape(1, 1)
+                    case 2 | 3:
+                        batch = arr[np.newaxis]
+                    case _:
+                        raise ValueError(f"Unsupported prediction shape: {arr.shape}")
+
+                n_channels = batch.shape[1]
+                output_tile_extent = batch.shape[2:] if batch.ndim > 2 else (1, 1)
+
+                if mask_builder is None:
+                    async with mask_builder_lock:
+                        if mask_builder is None:
+                            mask_builder = MaskBuilder(
+                                source_extents=(extent_y, extent_x),
+                                source_tile_extent=model_config["tile_size"],
+                                output_tile_extent=output_tile_extent,
+                                stride=stride,
+                                n_channels=n_channels,
+                                storage="memmap",
+                            )
+
+                mask_builder.update_batch(
+                    batch=batch,
+                    coords=np.array([[y, x]], dtype=np.int64),
+                )
 
             for x, y in grid_tiles(
                 slide_extent=(extent_x, extent_y),
@@ -100,11 +134,23 @@ async def process_tile(x: int, y: int) -> None:
 
             await asyncio.wait(tasks)
 
-        mask_builder.flush()
-        mask_builder.save(
+        if mask_builder is None:
+            raise RuntimeError("No tiles produced predictions; heatmap not created.")
+
+        result = mask_builder.finalize()
+        vips_image = mask_builder.resize_to_source(result)
+        vips_image = (vips_image * 255).cast(pyvips.BandFormat.UCHAR)
+        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+        vips_image.tiffsave(
             output_path,
-            tile_height=output_bigtiff_tile_height,
+            bigtiff=True,
+            compression=pyvips.enums.ForeignTiffCompression.DEFLATE,
+            tile=True,
             tile_width=output_bigtiff_tile_width,
+            tile_height=output_bigtiff_tile_height,
+            xres=1000 / mpp_x,
+            yres=1000 / mpp_y,
+            pyramid=True,
         )
         mask_builder.cleanup()
         return output_path
diff --git a/misc/tile_heatmap_builder.py b/misc/tile_heatmap_builder.py
deleted file mode 100644
index 843536c..0000000
--- a/misc/tile_heatmap_builder.py
+++ /dev/null
@@ -1,75 +0,0 @@
-import tempfile
-from pathlib import Path
-
-import numpy as np
-import pyvips
-
-
-class TileHeatmapBuilder:
-    def __init__(
-        self, extent_x: int, extent_y: int, mpp_x: float, mpp_y: float
-    ) -> None:
-        self.extent_x = extent_x
-        self.extent_y = extent_y
-        self.mpp_x = mpp_x
-        self.mpp_y = mpp_y
-
-        # Create temporary files
-        self.temp_dir = tempfile.TemporaryDirectory()
-        self.image_path = Path(self.temp_dir.name) / "image.dat"
-        self.count_path = Path(self.temp_dir.name) / "count.dat"
-
-        self.image = np.memmap(
-            str(self.image_path),
-            dtype=np.float32,
-            mode="w+",
-            shape=(self.extent_y, self.extent_x),
-        )
-
-        self.count = np.memmap(
-            str(self.count_path),
-            dtype=np.uint8,
-            mode="w+",
-            shape=(self.extent_y, self.extent_x),
-        )
-
-    def update(self, tile: np.ndarray, x: int, y: int) -> None:
-        mm_y, mm_x = self.image[y : y + tile.shape[0], x : x + tile.shape[1]].shape
-        self.image[y : y + mm_y, x : x + mm_x] += tile[:mm_y, :mm_x]
-        self.count[y : y + mm_y, x : x + mm_x] += 1
-
-    def flush(self) -> None:
-        self.image.flush()
-        self.count.flush()
-
-    def save(self, output_path: str, tile_width: int, tile_height: int) -> None:
-        image_vips = pyvips.Image.new_from_array(self.image)
-        count_vips = pyvips.Image.new_from_array(self.count)
-
-        image_vips /= count_vips
-        image_vips *= 255
-        image_vips = image_vips.cast(pyvips.BandFormat.UCHAR)
-
-        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
-        image_vips.tiffsave(
-            output_path,
-            bigtiff=True,
-            compression=pyvips.enums.ForeignTiffCompression.DEFLATE,
-            tile=True,
-            tile_width=tile_width,
-            tile_height=tile_height,
-            xres=1000 / self.mpp_x,
-            yres=1000 / self.mpp_y,
-            pyramid=True,
-        )
-
-    def cleanup(self) -> None:
-        if hasattr(self, "image"):
-            del self.image
-        if hasattr(self, "count"):
-            del self.count
-
-        self.temp_dir.cleanup()
-
-    def __del__(self) -> None:
-        self.cleanup()

From bc046bccef1b58a8ddf9d52f2507ef2c74171083 Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Thu, 14 May 2026 13:44:43 +0200
Subject: [PATCH 03/23] fix

Co-authored-by: Copilot <copilot@github.com>
---
 builders/heatmap_builder.py | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/builders/heatmap_builder.py b/builders/heatmap_builder.py
index d44a5ab..b390b0c 100644
--- a/builders/heatmap_builder.py
+++ b/builders/heatmap_builder.py
@@ -92,13 +92,7 @@ async def process_tile(x: int, y: int) -> None:
                 arr = np.asarray(prediction)
 
                 # Normalize to (B, C, *spatial) shape
-                match arr.ndim:
-                    case 0:
-                        batch = arr.reshape(1, 1)
-                    case 2 | 3:
-                        batch = arr[np.newaxis]
-                    case _:
-                        raise ValueError(f"Unsupported prediction shape: {arr.shape}")
+                batch = np.atleast_1d(arr)[np.newaxis, ...]
 
                 n_channels = batch.shape[1]
                 output_tile_extent = batch.shape[2:] if batch.ndim > 2 else (1, 1)

From b6fd99eacc21efb8d301fd5fc4a75f025c3e0ec5 Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Thu, 14 May 2026 14:46:19 +0200
Subject: [PATCH 04/23] feat: add foundation note

---
 docs/guides/adding-models.md | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/docs/guides/adding-models.md b/docs/guides/adding-models.md
index 440ec92..f6bed42 100644
--- a/docs/guides/adding-models.md
+++ b/docs/guides/adding-models.md
@@ -164,6 +164,39 @@ What each step does:
 - `.transpose(2, 0, 1)`: converts image to `CHW` layout expected by ONNX model.
 - `await self.predict(image)`: sends item to batching queue and waits for matching output.
 
+### Using Foundation Models from Your Model
+
+If you are deploying a model that is a downstream head (e.g., an MLP or Attention layer trained on top of a foundation model like Virchow2 or Prov-GigaPath), you **do not** need to re-export the entire foundation model into your ONNX artifact. 
+
+Because foundation models are already deployed as independent services within the cluster, your model can directly invoke them via Ray Serve handles. In your `root` or `predict` method, simply call the foundation model first, then pass its output to your custom layers.
+
+```python
+# In your __init__ or reconfigure method:
+from ray import serve
+self.foundation_model = serve.get_app_handle("virchow2")
+
+# In your request handler:
+@fastapi.post("/")
+async def root(self, request: Request):
+    # 1. Fetch raw image bytes from request
+    ...
+    
+    # 2. Call the foundation model via its Serve handle
+    # (assuming your custom model needs the raw `np.ndarray` image)
+    embedding = await self.foundation_model.predict.remote(image)
+    
+    # 3. Pass the embedding to your own model's predict endpoint
+    return await self.predict(embedding)
+```
+
+## Whole-Slide (WSI) Inference and Output Builders
+
+When predicting on an entire Whole-Slide Image (WSI):
+
+1. **Heatmaps:** If your model's WSI output should be a spatial heatmap (e.g., probability maps or segmentation masks overlaying the WSI), you **do not need to implement WSI logic**. The cluster already provides a universal `HeatmapBuilder` service (running under `/heatmap-builder`). Users can pass your model's ID to the heatmap builder via the SDK, and it will tile the image, aggregate all localized predictions seamlessly, and output a multi-resolution BigTIFF mask automatically.
+
+2. **Custom WSI Aggregations (Non-Heatmap Outputs):** If your model generates something else across the entire slide (for example, a single slide-level scalar score, diagnostic classification, custom tabular statistics, embedded feature bags), you must **implement your own WSI aggregator service**. You should create a custom Application (similar to `HeatmapBuilder`) that takes paths to WSI files, iterates through the WSI tiles querying your base model for each tile, and correctly aggregates the results into your desired slide-level output format.
+
 ### Application binding
 
 ```python

From a68986a418cae755fbb6deebc7c765a53bdeba0f Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Thu, 14 May 2026 15:17:33 +0200
Subject: [PATCH 05/23] fix

---
 docs/guides/adding-models.md | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/docs/guides/adding-models.md b/docs/guides/adding-models.md
index f6bed42..876b206 100644
--- a/docs/guides/adding-models.md
+++ b/docs/guides/adding-models.md
@@ -164,6 +164,14 @@ What each step does:
 - `.transpose(2, 0, 1)`: converts image to `CHW` layout expected by ONNX model.
 - `await self.predict(image)`: sends item to batching queue and waits for matching output.
 
+### Application binding
+
+```python
+app = MyModel.bind()
+```
+
+This exported symbol is what `import_path: models.my_model:app` points to in Helm.
+
 ### Using Foundation Models from Your Model
 
 If you are deploying a model that is a downstream head (e.g., an MLP or Attention layer trained on top of a foundation model like Virchow2 or Prov-GigaPath), you **do not** need to re-export the entire foundation model into your ONNX artifact. 
@@ -197,13 +205,6 @@ When predicting on an entire Whole-Slide Image (WSI):
 
 2. **Custom WSI Aggregations (Non-Heatmap Outputs):** If your model generates something else across the entire slide (for example, a single slide-level scalar score, diagnostic classification, custom tabular statistics, embedded feature bags), you must **implement your own WSI aggregator service**. You should create a custom Application (similar to `HeatmapBuilder`) that takes paths to WSI files, iterates through the WSI tiles querying your base model for each tile, and correctly aggregates the results into your desired slide-level output format.
 
-### Application binding
-
-```python
-app = MyModel.bind()
-```
-
-This exported symbol is what `import_path: models.my_model:app` points to in Helm.
 
 ## Next Step
 

From f2abdb07718f4127af1b48e65c3d4014412eceb1 Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Thu, 14 May 2026 15:24:35 +0200
Subject: [PATCH 06/23] fix

---
 builders/heatmap_builder.py | 35 ++++++++++++++++++++---------------
 1 file changed, 20 insertions(+), 15 deletions(-)

diff --git a/builders/heatmap_builder.py b/builders/heatmap_builder.py
index b390b0c..e3959c8 100644
--- a/builders/heatmap_builder.py
+++ b/builders/heatmap_builder.py
@@ -3,6 +3,8 @@
 from pathlib import Path
 from typing import Any, TypedDict
 
+import numpy as np
+import pyvips
 from fastapi import FastAPI
 from ray import serve
 
@@ -36,8 +38,6 @@ async def root(
         output_bigtiff_tile_height: int,
         output_bigtiff_tile_width: int,
     ) -> str:
-        import numpy as np
-        import pyvips
         from ratiopath.masks.mask_builders.mask_builder import MaskBuilder
         from ratiopath.openslide import OpenSlide
         from ratiopath.tiling import grid_tiles
@@ -50,6 +50,9 @@ async def root(
 
         loop = asyncio.get_running_loop()
         tasks: set[asyncio.Task[Any]] = set()
+        mask_builder: MaskBuilder | None = None
+        mask_builder_lock = asyncio.Lock()
+        update_lock = asyncio.Lock()
 
         with (
             OpenSlide(slide_path) as slide,
@@ -67,9 +70,6 @@ async def root(
             scale_x = tissue_extent_x / extent_x
             scale_y = tissue_extent_y / extent_y
 
-            mask_builder: MaskBuilder | None = None
-            mask_builder_lock = asyncio.Lock()
-
             async def process_tile(x: int, y: int) -> None:
                 nonlocal mask_builder
                 tile = await loop.run_in_executor(
@@ -89,13 +89,17 @@ async def process_tile(x: int, y: int) -> None:
                     return
 
                 prediction = await model.predict.remote(tile)
-                arr = np.asarray(prediction)
+                arr = np.asarray(prediction, dtype=np.float32)
 
-                # Normalize to (B, C, *spatial) shape
-                batch = np.atleast_1d(arr)[np.newaxis, ...]
+                if arr.ndim == 2:
+                    batch = arr[np.newaxis, np.newaxis, ...]
+                elif arr.ndim == 3:
+                    batch = arr[np.newaxis, ...]
+                else:
+                    raise ValueError(f"Unexpected prediction shape: {arr.shape}")
 
                 n_channels = batch.shape[1]
-                output_tile_extent = batch.shape[2:] if batch.ndim > 2 else (1, 1)
+                output_tile_extent = batch.shape[2:]
 
                 if mask_builder is None:
                     async with mask_builder_lock:
@@ -109,10 +113,11 @@ async def process_tile(x: int, y: int) -> None:
                                 storage="memmap",
                             )
 
-                mask_builder.update_batch(
-                    batch=batch,
-                    coords=np.array([[y, x]], dtype=np.int64),
-                )
+                async with update_lock:
+                    mask_builder.update_batch(
+                        batch=batch,
+                        coords=np.array([[y, x]], dtype=np.int64),
+                    )
 
             for x, y in grid_tiles(
                 slide_extent=(extent_x, extent_y),
@@ -123,7 +128,6 @@ async def process_tile(x: int, y: int) -> None:
                     _, tasks = await asyncio.wait(
                         tasks, return_when=asyncio.FIRST_COMPLETED
                     )
-
                 tasks.add(asyncio.create_task(process_tile(x, y)))
 
             await asyncio.wait(tasks)
@@ -131,7 +135,7 @@ async def process_tile(x: int, y: int) -> None:
         if mask_builder is None:
             raise RuntimeError("No tiles produced predictions; heatmap not created.")
 
-        result = mask_builder.finalize()
+        result = np.nan_to_num(mask_builder.finalize(), nan=0.0)
         vips_image = mask_builder.resize_to_source(result)
         vips_image = (vips_image * 255).cast(pyvips.BandFormat.UCHAR)
         Path(output_path).parent.mkdir(parents=True, exist_ok=True)
@@ -147,6 +151,7 @@ async def process_tile(x: int, y: int) -> None:
             pyramid=True,
         )
         mask_builder.cleanup()
+
         return output_path
 
 

From a8dda47dbdf1d0dbef0034662d2f0f85cd5f25df Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Thu, 14 May 2026 20:12:24 +0200
Subject: [PATCH 07/23] fix compiler

---
 builders/heatmap_builder.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/builders/heatmap_builder.py b/builders/heatmap_builder.py
index e3959c8..cf0ba69 100644
--- a/builders/heatmap_builder.py
+++ b/builders/heatmap_builder.py
@@ -112,6 +112,7 @@ async def process_tile(x: int, y: int) -> None:
                                 n_channels=n_channels,
                                 storage="memmap",
                             )
+                assert mask_builder is not None
 
                 async with update_lock:
                     mask_builder.update_batch(

From 2ca73c9f1725fcd696f82d5a8723f95694ab8825 Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Thu, 14 May 2026 20:19:03 +0200
Subject: [PATCH 08/23] README update

---
 README.md | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 568213c..4288949 100644
--- a/README.md
+++ b/README.md
@@ -4,15 +4,17 @@ Model deployment infrastructure for RationAI using Ray Serve on Kubernetes.
 
 This repository contains:
 
-- A Helm chart (`helm/rayservice/`) that renders and deploys a KubeRay `RayService`.
-- A static RayService manifest (`ray-service.yaml`) for reference/manual apply workflows.
-- Model implementations under `models/` (reference: `models/binary_classifier.py`).
-- Documentation under `docs/` (MkDocs).
+- `builders/`: WSI output aggregation services (e.g., `heatmap_builder.py`).
+- `docker/`: Dockerfiles for building CPU and GPU environments.
+- `docs/`: MkDocs documentation and architecture guides.
+- `helm/rayservice/`: A Helm chart that renders and deploys a KubeRay `RayService`.
+- `models/`: Python entrypoints for model implementations (e.g., `binary_classifier.py`, `virchow2.py`).
 
 ## Documentation
 
 - MkDocs content: `docs/`
 - Key pages:
+  - `docs/available-models.md`
   - `docs/get-started/quick-start.md`
   - `docs/guides/deployment-guide.md`
   - `docs/guides/adding-models.md`

From 7ccbc0706d85f5a3941e9a0004c10bac859cacf7 Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Thu, 14 May 2026 20:27:09 +0200
Subject: [PATCH 09/23] PR review comments

Co-authored-by: Copilot <copilot@github.com>
---
 builders/heatmap_builder.py | 56 +++++++++++++++++++++++--------------
 1 file changed, 35 insertions(+), 21 deletions(-)

diff --git a/builders/heatmap_builder.py b/builders/heatmap_builder.py
index cf0ba69..46a35f2 100644
--- a/builders/heatmap_builder.py
+++ b/builders/heatmap_builder.py
@@ -115,9 +115,12 @@ async def process_tile(x: int, y: int) -> None:
                 assert mask_builder is not None
 
                 async with update_lock:
-                    mask_builder.update_batch(
-                        batch=batch,
-                        coords=np.array([[y, x]], dtype=np.int64),
+                    await loop.run_in_executor(
+                        executor,
+                        lambda: mask_builder.update_batch(
+                            batch=batch,
+                            coords=np.array([[y, x]], dtype=np.int64),
+                        ),
                     )
 
             for x, y in grid_tiles(
@@ -126,32 +129,43 @@ async def process_tile(x: int, y: int) -> None:
                 stride=(stride, stride),
             ):
                 if len(tasks) >= self.max_concurrent_tasks:
-                    _, tasks = await asyncio.wait(
+                    done, tasks = await asyncio.wait(
                         tasks, return_when=asyncio.FIRST_COMPLETED
                     )
+                    for task in done:
+                        task.result()
                 tasks.add(asyncio.create_task(process_tile(x, y)))
 
-            await asyncio.wait(tasks)
+            if tasks:
+                done, _ = await asyncio.wait(tasks)
+                for task in done:
+                    task.result()
 
         if mask_builder is None:
             raise RuntimeError("No tiles produced predictions; heatmap not created.")
 
-        result = np.nan_to_num(mask_builder.finalize(), nan=0.0)
-        vips_image = mask_builder.resize_to_source(result)
-        vips_image = (vips_image * 255).cast(pyvips.BandFormat.UCHAR)
-        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
-        vips_image.tiffsave(
-            output_path,
-            bigtiff=True,
-            compression=pyvips.enums.ForeignTiffCompression.DEFLATE,
-            tile=True,
-            tile_width=output_bigtiff_tile_width,
-            tile_height=output_bigtiff_tile_height,
-            xres=1000 / mpp_x,
-            yres=1000 / mpp_y,
-            pyramid=True,
-        )
-        mask_builder.cleanup()
+        try:
+            result = np.nan_to_num(mask_builder.finalize(), nan=0.0, copy=False)
+            vips_image = mask_builder.resize_to_source(result)
+            vips_image = (vips_image * 255).cast(pyvips.BandFormat.UCHAR)
+            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+
+            await loop.run_in_executor(
+                None,
+                lambda: vips_image.tiffsave(
+                    output_path,
+                    bigtiff=True,
+                    compression=pyvips.enums.ForeignTiffCompression.DEFLATE,
+                    tile=True,
+                    tile_width=output_bigtiff_tile_width,
+                    tile_height=output_bigtiff_tile_height,
+                    xres=1000 / mpp_x,
+                    yres=1000 / mpp_y,
+                    pyramid=True,
+                ),
+            )
+        finally:
+            mask_builder.cleanup()
 
         return output_path
 

From e62899b1729a55cf06a6e521930a23c5cf3a1a8c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ji=C5=99=C3=AD=20=C5=A0t=C3=ADpek?=
 <91186480+Jurgee@users.noreply.github.com>
Date: Fri, 15 May 2026 14:56:53 +0200
Subject: [PATCH 10/23] Update builders/heatmap_builder.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Matěj Pekár <matej.pekar120@gmail.com>
---
 builders/heatmap_builder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/builders/heatmap_builder.py b/builders/heatmap_builder.py
index 46a35f2..bf94c83 100644
--- a/builders/heatmap_builder.py
+++ b/builders/heatmap_builder.py
@@ -38,7 +38,7 @@ async def root(
         output_bigtiff_tile_height: int,
         output_bigtiff_tile_width: int,
     ) -> str:
-        from ratiopath.masks.mask_builders.mask_builder import MaskBuilder
+        from ratiopath.masks.mask_builders import MaskBuilder
         from ratiopath.openslide import OpenSlide
         from ratiopath.tiling import grid_tiles
 

From 4bffeee1d5d96d64d4bf61c44b779c77a271c209 Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Fri, 15 May 2026 17:23:55 +0200
Subject: [PATCH 11/23] review fix

Co-authored-by: Copilot <copilot@github.com>
---
 builders/heatmap_builder.py                | 77 ++++++++--------------
 docs/guides/adding-models.md               | 16 +++++
 helm/rayservice/applications/episeg-1.yaml |  2 +
 models/semantic_segmentation.py            | 15 ++++-
 4 files changed, 59 insertions(+), 51 deletions(-)

diff --git a/builders/heatmap_builder.py b/builders/heatmap_builder.py
index bf94c83..6c65564 100644
--- a/builders/heatmap_builder.py
+++ b/builders/heatmap_builder.py
@@ -46,14 +46,13 @@ async def root(
 
         model = serve.get_app_handle(model_id)
         model_config = await model.get_config.remote()
-        stride: int = round(stride_fraction * model_config["tile_size"])
+        tile_size: int = model_config["tile_size"]
+        output_tile_size: int = model_config["output_tile_size"]
+        n_channels: int = model_config["n_channels"]
+        stride: int = round(stride_fraction * tile_size)
 
         loop = asyncio.get_running_loop()
         tasks: set[asyncio.Task[Any]] = set()
-        mask_builder: MaskBuilder | None = None
-        mask_builder_lock = asyncio.Lock()
-        update_lock = asyncio.Lock()
-
         with (
             OpenSlide(slide_path) as slide,
             OpenSlide(tissue_mask_path) as tissue_slide,
@@ -69,9 +68,16 @@ async def root(
             ]
             scale_x = tissue_extent_x / extent_x
             scale_y = tissue_extent_y / extent_y
+            mask_builder = MaskBuilder(
+                source_extents=(extent_y, extent_x),
+                source_tile_extent=tile_size,
+                output_tile_extent=output_tile_size,
+                stride=stride,
+                n_channels=n_channels,
+                storage="memmap",
+            )
 
             async def process_tile(x: int, y: int) -> None:
-                nonlocal mask_builder
                 tile = await loop.run_in_executor(
                     executor,
                     fetch_tissue_tile,
@@ -83,7 +89,7 @@ async def process_tile(x: int, y: int) -> None:
                     scale_x,
                     scale_y,
                     tissue_level,
-                    model_config["tile_size"],
+                    tile_size,
                 )
                 if tile is None:
                     return
@@ -98,34 +104,14 @@ async def process_tile(x: int, y: int) -> None:
                 else:
                     raise ValueError(f"Unexpected prediction shape: {arr.shape}")
 
-                n_channels = batch.shape[1]
-                output_tile_extent = batch.shape[2:]
-
-                if mask_builder is None:
-                    async with mask_builder_lock:
-                        if mask_builder is None:
-                            mask_builder = MaskBuilder(
-                                source_extents=(extent_y, extent_x),
-                                source_tile_extent=model_config["tile_size"],
-                                output_tile_extent=output_tile_extent,
-                                stride=stride,
-                                n_channels=n_channels,
-                                storage="memmap",
-                            )
-                assert mask_builder is not None
-
-                async with update_lock:
-                    await loop.run_in_executor(
-                        executor,
-                        lambda: mask_builder.update_batch(
-                            batch=batch,
-                            coords=np.array([[y, x]], dtype=np.int64),
-                        ),
-                    )
+                mask_builder.update_batch(
+                    batch=batch,
+                    coords=np.array([[y, x]], dtype=np.int64),
+                )
 
             for x, y in grid_tiles(
                 slide_extent=(extent_x, extent_y),
-                tile_extent=(model_config["tile_size"], model_config["tile_size"]),
+                tile_extent=(tile_size, tile_size),
                 stride=(stride, stride),
             ):
                 if len(tasks) >= self.max_concurrent_tasks:
@@ -141,28 +127,21 @@ async def process_tile(x: int, y: int) -> None:
                 for task in done:
                     task.result()
 
-        if mask_builder is None:
-            raise RuntimeError("No tiles produced predictions; heatmap not created.")
-
         try:
             result = np.nan_to_num(mask_builder.finalize(), nan=0.0, copy=False)
             vips_image = mask_builder.resize_to_source(result)
             vips_image = (vips_image * 255).cast(pyvips.BandFormat.UCHAR)
             Path(output_path).parent.mkdir(parents=True, exist_ok=True)
-
-            await loop.run_in_executor(
-                None,
-                lambda: vips_image.tiffsave(
-                    output_path,
-                    bigtiff=True,
-                    compression=pyvips.enums.ForeignTiffCompression.DEFLATE,
-                    tile=True,
-                    tile_width=output_bigtiff_tile_width,
-                    tile_height=output_bigtiff_tile_height,
-                    xres=1000 / mpp_x,
-                    yres=1000 / mpp_y,
-                    pyramid=True,
-                ),
+            vips_image.tiffsave(
+                output_path,
+                bigtiff=True,
+                compression=pyvips.enums.ForeignTiffCompression.DEFLATE,
+                tile=True,
+                tile_width=output_bigtiff_tile_width,
+                tile_height=output_bigtiff_tile_height,
+                xres=1000 / mpp_x,
+                yres=1000 / mpp_y,
+                pyramid=True,
             )
         finally:
             mask_builder.cleanup()
diff --git a/docs/guides/adding-models.md b/docs/guides/adding-models.md
index 876b206..21a32d1 100644
--- a/docs/guides/adding-models.md
+++ b/docs/guides/adding-models.md
@@ -141,6 +141,22 @@ What happens with data:
 - Output tensor is flattened and returned as a Python list.
 - Ray Serve maps each list item back to the original HTTP request.
 
+### `get_config`: expose model settings for builders
+
+If your model is used by any Whole-Slide Inference builder (for example `HeatmapBuilder`), you must provide a `get_config` method that builders can call through a Serve handle. The builder uses this to read `tile_size`, `output_tile_size`, `n_channels`, and `mpp` so it can pick the right tiling grid and resolution.
+
+```python
+async def get_config(self) -> dict[str, Any]:
+  return {
+    "tile_size": self.tile_size,
+    "output_tile_size": self.output_tile_size,
+    "n_channels": self.n_channels,
+    "mpp": self.mpp,
+  }
+```
+
+The builder calls it with `await model.get_config.remote()`; keep it cheap and avoid any I/O.
+
 ### `root`: HTTP request parsing and serialization
 
 ```python
diff --git a/helm/rayservice/applications/episeg-1.yaml b/helm/rayservice/applications/episeg-1.yaml
index 5c8e9db..66a08af 100644
--- a/helm/rayservice/applications/episeg-1.yaml
+++ b/helm/rayservice/applications/episeg-1.yaml
@@ -20,6 +20,8 @@
             MLFLOW_TRACKING_URI: http://mlflow.rationai-mlflow:5000
       user_config:
         tile_size: 1024
+        output_tile_size: 1024
+        n_channels: 1
         mpp: 0.468
         max_batch_size: 8
         batch_wait_timeout_s: 0.1
diff --git a/models/semantic_segmentation.py b/models/semantic_segmentation.py
index cc1611b..67242c7 100644
--- a/models/semantic_segmentation.py
+++ b/models/semantic_segmentation.py
@@ -9,6 +9,8 @@
 
 class Config(TypedDict):
     tile_size: int
+    output_tile_size: int
+    n_channels: int
     mpp: float
     model: dict[str, Any]
     max_batch_size: int
@@ -29,6 +31,8 @@ class SemanticSegmentation:
     """Semantic segmentation for tissue tiles using ONNX Runtime with GPU and TensorRT support."""
 
     tile_size: int
+    output_tile_size: int
+    n_channels: int
 
     def __init__(self) -> None:
         import lz4.frame
@@ -42,6 +46,8 @@ def reconfigure(self, config: Config) -> None:
         import onnxruntime as ort
 
         self.tile_size = config["tile_size"]
+        self.output_tile_size = config["output_tile_size"]
+        self.n_channels = config["n_channels"]
         self.mpp = config["mpp"]
 
         cache_path = config["trt_cache_path"]
@@ -116,8 +122,13 @@ def reconfigure(self, config: Config) -> None:
         self.predict.set_batch_wait_timeout_s(config["batch_wait_timeout_s"])  # type: ignore[attr-defined]
 
     def get_config(self) -> dict[str, Any]:
-        """Return the current configuration (tile size and mpp)."""
-        return {"tile_size": self.tile_size, "mpp": self.mpp}
+        """Return the current configuration for builders."""
+        return {
+            "tile_size": self.tile_size,
+            "output_tile_size": self.output_tile_size,
+            "n_channels": self.n_channels,
+            "mpp": self.mpp,
+        }
 
     @serve.batch
     async def predict(

From 396af101f81d18d8f08d573d71be6577ce70eb7c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ji=C5=99=C3=AD=20=C5=A0t=C3=ADpek?=
 <91186480+Jurgee@users.noreply.github.com>
Date: Fri, 15 May 2026 18:26:29 +0200
Subject: [PATCH 12/23] Potential fix for pull request finding

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
---
 docs/guides/adding-models.md | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/docs/guides/adding-models.md b/docs/guides/adding-models.md
index 21a32d1..e736e3b 100644
--- a/docs/guides/adding-models.md
+++ b/docs/guides/adding-models.md
@@ -192,24 +192,29 @@ This exported symbol is what `import_path: models.my_model:app` points to in Hel
 
 If you are deploying a model that is a downstream head (e.g., an MLP or Attention layer trained on top of a foundation model like Virchow2 or Prov-GigaPath), you **do not** need to re-export the entire foundation model into your ONNX artifact. 
 
-Because foundation models are already deployed as independent services within the cluster, your model can directly invoke them via Ray Serve handles. In your `root` or `predict` method, simply call the foundation model first, then pass its output to your custom layers.
+Because foundation models are already deployed as independent services within the cluster, your model can directly invoke them via Ray Serve handles. In your `root` or `predict` method, call the foundation model first, then pass its output to your custom layers. If you call a foundation model's `predict` method directly, do **not** pass a raw `np.ndarray` image; first apply the same preprocessing/transforms that deployment expects, or call the model's ingress/request path instead.
 
 ```python
 # In your __init__ or reconfigure method:
 from ray import serve
 self.foundation_model = serve.get_app_handle("virchow2")
+self.foundation_transform = build_virchow2_transform()
 
 # In your request handler:
 @fastapi.post("/")
 async def root(self, request: Request):
-    # 1. Fetch raw image bytes from request
+    # 1. Fetch raw image bytes from request and decode to an image / np.ndarray
     ...
+
+    # 2. Apply the same transforms used by the foundation model deployment
+    image_tensor = self.foundation_transform(image)
+    if image_tensor.ndim == 3:
+        image_tensor = image_tensor.unsqueeze(0)
+
+    # 3. Call the foundation model with the transformed tensor batch
+    embedding = await self.foundation_model.predict.remote(image_tensor)
     
-    # 2. Call the foundation model via its Serve handle
-    # (assuming your custom model needs the raw `np.ndarray` image)
-    embedding = await self.foundation_model.predict.remote(image)
-    
-    # 3. Pass the embedding to your own model's predict endpoint
+    # 4. Pass the embedding to your own model's predict endpoint
     return await self.predict(embedding)
 ```
 

From 25262459e97ac197cc5b6ee126f7d11633436a8c Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Fri, 15 May 2026 18:28:20 +0200
Subject: [PATCH 13/23] fix: try block move

Co-authored-by: Copilot <copilot@github.com>
---
 builders/heatmap_builder.py | 132 ++++++++++++++++++------------------
 1 file changed, 66 insertions(+), 66 deletions(-)

diff --git a/builders/heatmap_builder.py b/builders/heatmap_builder.py
index 6c65564..bae6520 100644
--- a/builders/heatmap_builder.py
+++ b/builders/heatmap_builder.py
@@ -76,75 +76,75 @@ async def root(
                 n_channels=n_channels,
                 storage="memmap",
             )
-
-            async def process_tile(x: int, y: int) -> None:
-                tile = await loop.run_in_executor(
-                    executor,
-                    fetch_tissue_tile,
-                    slide,
-                    tissue_slide,
-                    x,
-                    y,
-                    level,
-                    scale_x,
-                    scale_y,
-                    tissue_level,
-                    tile_size,
-                )
-                if tile is None:
-                    return
-
-                prediction = await model.predict.remote(tile)
-                arr = np.asarray(prediction, dtype=np.float32)
-
-                if arr.ndim == 2:
-                    batch = arr[np.newaxis, np.newaxis, ...]
-                elif arr.ndim == 3:
-                    batch = arr[np.newaxis, ...]
-                else:
-                    raise ValueError(f"Unexpected prediction shape: {arr.shape}")
-
-                mask_builder.update_batch(
-                    batch=batch,
-                    coords=np.array([[y, x]], dtype=np.int64),
-                )
-
-            for x, y in grid_tiles(
-                slide_extent=(extent_x, extent_y),
-                tile_extent=(tile_size, tile_size),
-                stride=(stride, stride),
-            ):
-                if len(tasks) >= self.max_concurrent_tasks:
-                    done, tasks = await asyncio.wait(
-                        tasks, return_when=asyncio.FIRST_COMPLETED
+            try:
+
+                async def process_tile(x: int, y: int) -> None:
+                    tile = await loop.run_in_executor(
+                        executor,
+                        fetch_tissue_tile,
+                        slide,
+                        tissue_slide,
+                        x,
+                        y,
+                        level,
+                        scale_x,
+                        scale_y,
+                        tissue_level,
+                        tile_size,
+                    )
+                    if tile is None:
+                        return
+
+                    prediction = await model.predict.remote(tile)
+                    arr = np.asarray(prediction, dtype=np.float32)
+
+                    if arr.ndim == 2:
+                        batch = arr[np.newaxis, np.newaxis, ...]
+                    elif arr.ndim == 3:
+                        batch = arr[np.newaxis, ...]
+                    else:
+                        raise ValueError(f"Unexpected prediction shape: {arr.shape}")
+
+                    mask_builder.update_batch(
+                        batch=batch,
+                        coords=np.array([[y, x]], dtype=np.int64),
                     )
+
+                for x, y in grid_tiles(
+                    slide_extent=(extent_x, extent_y),
+                    tile_extent=(tile_size, tile_size),
+                    stride=(stride, stride),
+                ):
+                    if len(tasks) >= self.max_concurrent_tasks:
+                        done, tasks = await asyncio.wait(
+                            tasks, return_when=asyncio.FIRST_COMPLETED
+                        )
+                        for task in done:
+                            task.result()
+                    tasks.add(asyncio.create_task(process_tile(x, y)))
+
+                if tasks:
+                    done, _ = await asyncio.wait(tasks)
                     for task in done:
                         task.result()
-                tasks.add(asyncio.create_task(process_tile(x, y)))
-
-            if tasks:
-                done, _ = await asyncio.wait(tasks)
-                for task in done:
-                    task.result()
-
-        try:
-            result = np.nan_to_num(mask_builder.finalize(), nan=0.0, copy=False)
-            vips_image = mask_builder.resize_to_source(result)
-            vips_image = (vips_image * 255).cast(pyvips.BandFormat.UCHAR)
-            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
-            vips_image.tiffsave(
-                output_path,
-                bigtiff=True,
-                compression=pyvips.enums.ForeignTiffCompression.DEFLATE,
-                tile=True,
-                tile_width=output_bigtiff_tile_width,
-                tile_height=output_bigtiff_tile_height,
-                xres=1000 / mpp_x,
-                yres=1000 / mpp_y,
-                pyramid=True,
-            )
-        finally:
-            mask_builder.cleanup()
+
+                result = np.nan_to_num(mask_builder.finalize(), nan=0.0, copy=False)
+                vips_image = mask_builder.resize_to_source(result)
+                vips_image = (vips_image * 255).cast(pyvips.BandFormat.UCHAR)
+                Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+                vips_image.tiffsave(
+                    output_path,
+                    bigtiff=True,
+                    compression=pyvips.enums.ForeignTiffCompression.DEFLATE,
+                    tile=True,
+                    tile_width=output_bigtiff_tile_width,
+                    tile_height=output_bigtiff_tile_height,
+                    xres=1000 / mpp_x,
+                    yres=1000 / mpp_y,
+                    pyramid=True,
+                )
+            finally:
+                mask_builder.cleanup()
 
         return output_path
 

From 4515c9ab7b2587f6cc3310b71f1d62c6e0e8866b Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Sun, 17 May 2026 15:02:10 +0200
Subject: [PATCH 14/23] new docker file

Co-authored-by: Copilot <copilot@github.com>
---
 docker/Dockerfile.cpu | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu
index 9f778f5..56397ef 100644
--- a/docker/Dockerfile.cpu
+++ b/docker/Dockerfile.cpu
@@ -52,4 +52,4 @@ RUN sudo apt-get update && sudo apt-get -y upgrade && \
 RUN sudo apt-get remove -y --purge systemd systemd-sysv && sudo apt-get autoremove --purge -y && sudo apt-get clean && sudo rm -rf /var/lib/apt/lists/* 
 
 RUN pip install --no-cache-dir \
-    onnxruntime lz4 ratiopath "mlflow<3.0"
\ No newline at end of file
+    onnxruntime lz4 ratiopath pyvips "mlflow<3.0"
\ No newline at end of file

From 294a99547dc121ac0bccf891a2043ff31121a1fa Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Sun, 17 May 2026 15:32:20 +0200
Subject: [PATCH 15/23] move import

Co-authored-by: Copilot <copilot@github.com>
---
 builders/heatmap_builder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/builders/heatmap_builder.py b/builders/heatmap_builder.py
index bae6520..f07d3bb 100644
--- a/builders/heatmap_builder.py
+++ b/builders/heatmap_builder.py
@@ -4,7 +4,6 @@
 from typing import Any, TypedDict
 
 import numpy as np
-import pyvips
 from fastapi import FastAPI
 from ray import serve
 
@@ -38,6 +37,7 @@ async def root(
         output_bigtiff_tile_height: int,
         output_bigtiff_tile_width: int,
     ) -> str:
+        import pyvips
         from ratiopath.masks.mask_builders import MaskBuilder
         from ratiopath.openslide import OpenSlide
         from ratiopath.tiling import grid_tiles

From 6fe42ee0d370cfc2a81e8a301753578acd80bc03 Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Sun, 17 May 2026 15:52:58 +0200
Subject: [PATCH 16/23] add transpose

---
 builders/heatmap_builder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/builders/heatmap_builder.py b/builders/heatmap_builder.py
index f07d3bb..cf2aafd 100644
--- a/builders/heatmap_builder.py
+++ b/builders/heatmap_builder.py
@@ -95,7 +95,7 @@ async def process_tile(x: int, y: int) -> None:
                     if tile is None:
                         return
 
-                    prediction = await model.predict.remote(tile)
+                    prediction = await model.predict.remote(tile.transpose(2, 0, 1))
                     arr = np.asarray(prediction, dtype=np.float32)
 
                     if arr.ndim == 2:

From c9ecf90e7163aba06d5a2c095b3be1476d23d71b Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Sun, 17 May 2026 16:48:13 +0200
Subject: [PATCH 17/23] remove transpose

Co-authored-by: Copilot <copilot@github.com>
---
 builders/heatmap_builder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/builders/heatmap_builder.py b/builders/heatmap_builder.py
index cf2aafd..f07d3bb 100644
--- a/builders/heatmap_builder.py
+++ b/builders/heatmap_builder.py
@@ -95,7 +95,7 @@ async def process_tile(x: int, y: int) -> None:
                     if tile is None:
                         return
 
-                    prediction = await model.predict.remote(tile.transpose(2, 0, 1))
+                    prediction = await model.predict.remote(tile)
                     arr = np.asarray(prediction, dtype=np.float32)
 
                     if arr.ndim == 2:

From 0285bb552a7a2725914b95ceac4c0d71ff05caf4 Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Sun, 17 May 2026 16:57:51 +0200
Subject: [PATCH 18/23] fix

Co-authored-by: Copilot <copilot@github.com>
---
 builders/heatmap_builder.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/builders/heatmap_builder.py b/builders/heatmap_builder.py
index f07d3bb..d02c6d7 100644
--- a/builders/heatmap_builder.py
+++ b/builders/heatmap_builder.py
@@ -128,7 +128,8 @@ async def process_tile(x: int, y: int) -> None:
                     for task in done:
                         task.result()
 
-                result = np.nan_to_num(mask_builder.finalize(), nan=0.0, copy=False)
+                result = np.asarray(mask_builder.finalize())
+                np.nan_to_num(result, nan=0.0, copy=False)
                 vips_image = mask_builder.resize_to_source(result)
                 vips_image = (vips_image * 255).cast(pyvips.BandFormat.UCHAR)
                 Path(output_path).parent.mkdir(parents=True, exist_ok=True)

From 9fcd84e9513206b89b2ba9cb11c110ea447f53b6 Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Sun, 17 May 2026 17:26:12 +0200
Subject: [PATCH 19/23] fix

Co-authored-by: Copilot <copilot@github.com>
---
 builders/heatmap_builder.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/builders/heatmap_builder.py b/builders/heatmap_builder.py
index d02c6d7..ec9ab98 100644
--- a/builders/heatmap_builder.py
+++ b/builders/heatmap_builder.py
@@ -129,7 +129,12 @@ async def process_tile(x: int, y: int) -> None:
                         task.result()
 
                 result = np.asarray(mask_builder.finalize())
-                np.nan_to_num(result, nan=0.0, copy=False)
+
+                if result.ndim == 2:
+                    result = result[np.newaxis, ...]
+                else:
+                    raise ValueError(f"Unexpected shape: {result.shape}")
+
                 vips_image = mask_builder.resize_to_source(result)
                 vips_image = (vips_image * 255).cast(pyvips.BandFormat.UCHAR)
                 Path(output_path).parent.mkdir(parents=True, exist_ok=True)

From dc2009e9247d12ac920851313acad2b0e2603460 Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Sun, 17 May 2026 17:37:09 +0200
Subject: [PATCH 20/23] add mask to finalize

Co-authored-by: Copilot <copilot@github.com>
---
 builders/heatmap_builder.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/builders/heatmap_builder.py b/builders/heatmap_builder.py
index ec9ab98..328403f 100644
--- a/builders/heatmap_builder.py
+++ b/builders/heatmap_builder.py
@@ -128,12 +128,10 @@ async def process_tile(x: int, y: int) -> None:
                     for task in done:
                         task.result()
 
-                result = np.asarray(mask_builder.finalize())
+                result = np.asarray(mask_builder.finalize()["mask"])
 
                 if result.ndim == 2:
                     result = result[np.newaxis, ...]
-                else:
-                    raise ValueError(f"Unexpected shape: {result.shape}")
 
                 vips_image = mask_builder.resize_to_source(result)
                 vips_image = (vips_image * 255).cast(pyvips.BandFormat.UCHAR)

From 6eed8acf2233536365f12460cf381066ad95510e Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Sun, 17 May 2026 18:05:44 +0200
Subject: [PATCH 21/23] new Docker image for cpu

---
 helm/rayservice/workers/cpu-workers.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/helm/rayservice/workers/cpu-workers.yaml b/helm/rayservice/workers/cpu-workers.yaml
index 0c98d85..eaa044b 100644
--- a/helm/rayservice/workers/cpu-workers.yaml
+++ b/helm/rayservice/workers/cpu-workers.yaml
@@ -11,7 +11,7 @@ template:
         type: RuntimeDefault
     containers:
       - name: ray-worker
-        image: cerit.io/rationai/model-service:2.54.0
+        image: cerit.io/rationai/model-service:2.55.0
         imagePullPolicy: Always
         resources:
           limits:

From 109c4909126d3098c33f23b3eb7b54fd9a3825d8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ji=C5=99=C3=AD=20=C5=A0t=C3=ADpek?=
 <91186480+Jurgee@users.noreply.github.com>
Date: Sun, 17 May 2026 19:30:52 +0200
Subject: [PATCH 22/23] Update builders/heatmap_builder.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Matěj Pekár <matej.pekar120@gmail.com>
---
 builders/heatmap_builder.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/builders/heatmap_builder.py b/builders/heatmap_builder.py
index 328403f..a1b7492 100644
--- a/builders/heatmap_builder.py
+++ b/builders/heatmap_builder.py
@@ -130,8 +130,6 @@ async def process_tile(x: int, y: int) -> None:
 
                 result = np.asarray(mask_builder.finalize()["mask"])
 
-                if result.ndim == 2:
-                    result = result[np.newaxis, ...]
 
                 vips_image = mask_builder.resize_to_source(result)
                 vips_image = (vips_image * 255).cast(pyvips.BandFormat.UCHAR)

From bedd9e4167bbace5689c5af1e8ec15937310b3d6 Mon Sep 17 00:00:00 2001
From: JiriStipek <567776@mail.muni.cz>
Date: Sun, 17 May 2026 19:33:24 +0200
Subject: [PATCH 23/23] lint fix

---
 builders/heatmap_builder.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/builders/heatmap_builder.py b/builders/heatmap_builder.py
index a1b7492..8b89d31 100644
--- a/builders/heatmap_builder.py
+++ b/builders/heatmap_builder.py
@@ -130,7 +130,6 @@ async def process_tile(x: int, y: int) -> None:
 
                 result = np.asarray(mask_builder.finalize()["mask"])
 
-
                 vips_image = mask_builder.resize_to_source(result)
                 vips_image = (vips_image * 255).cast(pyvips.BandFormat.UCHAR)
                 Path(output_path).parent.mkdir(parents=True, exist_ok=True)