Load tiles in parallel on workers and add options to TissueDetectionHE

pathml/core/slide_backends.py

@@ -5,6 +5,7 @@
 
 from io import BytesIO
 
+import dask
 import numpy as np
 import openslide
 from javabridge.jutil import JavaException
@@ -62,8 +63,14 @@ class OpenSlideBackend(SlideBackend):
  def __init__(self, filename):
  logger.info(f"OpenSlideBackend loading file at: {filename}")
  self.filename = filename
- self.slide = openslide.open_slide(filename=filename)
- self.level_count = self.slide.level_count
+
+ @property
+ def slide(self):
+ return openslide.open_slide(filename=self.filename)
+
+ @property
+ def level_count(self):
+ return self.slide.level_count
 
  def __repr__(self):
  return f"OpenSlideBackend('{self.filename}')"
@@ -211,9 +218,10 @@ def generate_tiles(self, shape=3000, stride=None, pad=False, level=0):
  for ix_i in range(n_tiles_i):
  for ix_j in range(n_tiles_j):
  coords = (int(ix_i * stride_i), int(ix_j * stride_j))
- # get image for tile
- tile_im = self.extract_region(location=coords, size=shape, level=level)
- yield pathml.core.tile.Tile(image=tile_im, coords=coords)
+ image = dask.delayed(self.extract_region)(
+ location=coords, size=shape, level=level
+ )
+ yield pathml.core.tile.Tile(image, coords=coords)
 
 
 def _init_logger():
@@ -421,7 +429,10 @@ def extract_region(
  f"Multi-level images not supported with series_as_channels=True. Input 'level={level}' invalid. Use 'level=0'."
  )
 
- javabridge.start_vm(class_path=bioformats.JARS, max_heap_size="100G")
+ javabridge.start_vm(
+ class_path=bioformats.JARS, max_heap_size="100G", run_headless=True
+ )
+
  with bioformats.ImageReader(str(self.filename), perform_init=True) as reader:
  # expand size
  logger.info(f"extracting region with input size = {size}")
@@ -593,27 +604,28 @@ def generate_tiles(self, shape=3000, stride=None, pad=False, level=0, **kwargs):
  for ix_j in range(n_tiles_j):
  coords = (int(ix_i * stride_i), int(ix_j * stride_j))
  if coords[0] + shape[0] < i and coords[1] + shape[1] < j:
- # get image for tile
- tile_im = self.extract_region(
+ image = dask.delayed(self.extract_region)(
  location=coords, size=shape, level=level, **kwargs
  )
-  yield pathml.core.tile.Tile(image=tile_im, coords=coords)
+ # Image on edge and needs to be padded with 0s
  else:
- unpaddedshape = (
+ unpadded_shape = (
  i - coords[0] if coords[0] + shape[0] > i else shape[0],
  j - coords[1] if coords[1] + shape[1] > j else shape[1],
  )
- tile_im = self.extract_region(
- location=coords, size=unpaddedshape, level=level, **kwargs
+ edge_image = dask.delayed(self.extract_region)(
+ location=coords, size=unpadded_shape, level=level, **kwargs
  )
- zeroarrayshape = list(tile_im.shape)
- zeroarrayshape[0], zeroarrayshape[1] = (
- list(shape)[0],
- list(shape)[1],
- )
- padded_im = np.zeros(zeroarrayshape)
- padded_im[: tile_im.shape[0], : tile_im.shape[1], ...] = tile_im
- yield pathml.core.tile.Tile(image=padded_im, coords=coords)
+
+ def pad(image):
+ """Pads edge tiles with zeros."""
+ padded = np.zeros((*shape, *image.shape[:-2]))
+ padded[: image.shape[0], : image.shape[1]] = image
+ return padded
+
+ # Need to delay to use shape of edge_image
+ image = dask.delayed(pad)(edge_image)
+ yield pathml.core.tile.Tile(image=image, coords=coords)
 
 
 class DICOMBackend(SlideBackend):
@@ -653,19 +665,25 @@ def __init__(self, filename):
  f"DICOM metadata: frame_shape={self.frame_shape}, nrows = {self.n_rows}, ncols = {self.n_cols}"
  )
 
- # actual file
- self.fp = DicomFile(self.filename, mode="rb")
- self.fp.is_little_endian = self.transfer_syntax_uid.is_little_endian
- self.fp.is_implicit_VR = self.transfer_syntax_uid.is_implicit_VR
+ fp = self.fp
+
  # need to do this to advance the file to the correct point, at the beginning of the pixels
- self.metadata = dcmread(self.fp, stop_before_pixels=True)
- self.pixel_data_offset = self.fp.tell()
- self.fp.seek(self.pixel_data_offset, 0)
+ self.metadata = dcmread(fp, stop_before_pixels=True)
+ pixel_data_offset = fp.tell()
+ fp.seek(pixel_data_offset, 0)
  # note that reading this tag is necessary to advance the file to correct position
- _ = TupleTag(self.fp.read_tag())
+ _ = TupleTag(fp.read_tag())
  # get basic offset table, to enable reading individual frames without loading entire image
- self.bot = self.get_bot(self.fp)
- self.first_frame = self.fp.tell()
+ self.bot = self.get_bot(fp)
+ self.first_frame = fp.tell()
+
+ @property
+ def fp(self):
+ """actual file"""
+ fp = DicomFile(self.filename, mode="rb")
+ fp.is_little_endian = self.transfer_syntax_uid.is_little_endian
+ fp.is_implicit_VR = self.transfer_syntax_uid.is_implicit_VR
+ return fp
 
  def __repr__(self):
  out = f"DICOMBackend('{self.filename}')\n"
@@ -807,7 +825,9 @@ def _read_frame(self, frame_ix):
  np.ndarray: pixel data of that frame
  """
  frame_offset = self.bot[int(frame_ix)]
- self.fp.seek(self.first_frame + frame_offset, 0)
+ # self.fp refers to a different filelike object each time it is accessed
+ fp = self.fp
+ fp.seek(self.first_frame + frame_offset, 0)
  try:
  stop_at = self.bot[frame_ix + 1] - frame_offset
  except IndexError:
@@ -816,11 +836,11 @@ def _read_frame(self, frame_ix):
  # A frame may comprised of multiple chunks
  chunks = []
  while True:
- tag = TupleTag(self.fp.read_tag())
+ tag = TupleTag(fp.read_tag())
  if n == stop_at or int(tag) == SequenceDelimiterTag:
  break
- length = self.fp.read_UL()
- chunks.append(self.fp.read(length))
+ length = fp.read_UL()
+ chunks.append(fp.read(length))
  n += 8 + length
 
  frame_bytes = b"".join(chunks)
@@ -899,7 +919,7 @@ def generate_tiles(self, shape, stride, pad, level=0, **kwargs):
  if i >= (self.n_frames - self.n_cols):
  continue
 
- frame_im = self.extract_region(location=i)
+ im = dask.delayed(self.extract_region)(location=i)
  coords = self._index_to_coords(i)
- frame_tile = pathml.core.tile.Tile(image=frame_im, coords=coords)
- yield frame_tile
+ tile = pathml.core.tile.Tile(image=im, coords=coords)
+ yield tile

pathml/core/slide_data.py

@@ -17,6 +17,7 @@
 import pathml.core
 import pathml.preprocessing.pipeline
 from pathml.core.slide_types import SlideType
+from pathml.preprocessing.transforms import DropTileException
 
 
 def infer_backend(path):
@@ -309,31 +310,47 @@ def run(
  )
 
  # map pipeline application onto each tile
- processed_tile_futures = []
+ futures = [
+ client.submit(pipeline.apply, tile)
+ for tile in self.generate_tiles(
+ level=level,
+ shape=tile_size,
+ stride=tile_stride,
+ pad=tile_pad,
+ **kwargs,
+ )
+ ]
 
- for tile in self.generate_tiles(
- level=level,
- shape=tile_size,
- stride=tile_stride,
- pad=tile_pad,
- **kwargs,
- ):
- if not tile.slide_type:
- tile.slide_type = self.slide_type
- # explicitly scatter data, i.e. send the tile data out to the cluster before applying the pipeline
- # according to dask, this can reduce scheduler burden and keep data on workers
- big_future = client.scatter(tile)
- f = client.submit(pipeline.apply, big_future)
- processed_tile_futures.append(f)
-
- # as tiles are processed, add them to h5
- for future, tile in dask.distributed.as_completed(
- processed_tile_futures, with_results=True
+ # After a worker processes a tile, add the tile to h5
+ for future, result in dask.distributed.as_completed(
+ futures, with_results=True, raise_errors=False
  ):
- self.tiles.add(tile)
+ if future.status == "finished":
+ self.tiles.add(result)
+ if future.status == "error":
+ typ, exc, tb = result
+ if typ is DropTileException:
+ pass
+ else:
+ raise exc.with_traceback(tb)
+ # TODO: Free memory used for tile
+ # Each in-memory future holding a Tile shows a size of 48 bytes on the Dask dashboard
+ # which clearly does not include image data.
+ # Could it be that loaded image data is somehow not being garbage collected with Tiles?
+
+ # # all of these still leave unmanaged memory on each worker
+ # future.release()
+ # future.cancel()
+ # del result
+ # del future
+ # del futures
 
  if shutdown_after:
  client.shutdown()
+ else:
+ pass
+ # Stopgap to free unmanaged memory on client before processing another slide
+ client.restart()
 
  else:
  for tile in self.generate_tiles(
@@ -343,8 +360,6 @@ def run(
  pad=tile_pad,
  **kwargs,
  ):
- if not tile.slide_type:
- tile.slide_type = self.slide_type
  pipeline.apply(tile)
  self.tiles.add(tile)
 
@@ -410,14 +425,19 @@ def generate_tiles(self, shape=3000, stride=None, pad=False, **kwargs):
  pathml.core.tile.Tile: Extracted Tile object
  """
  for tile in self.slide.generate_tiles(shape, stride, pad, **kwargs):
+ # TODO: move to worker!! (forces loading data on main thread)
+
  # add masks for tile, if possible
  # i.e. if the SlideData has a Masks object, and the tile has coordinates
  if self.masks is not None and tile.coords is not None:
  # masks not supported if pad=True
  # to implement, need to update Mask.slice to support slices that go beyond the full mask
  if not pad:
  i, j = tile.coords
- di, dj = tile.image.shape[0:2]
+ # Accessing image loads data on main thread
+ # dask.delayed waits until compute is called on worker
+ shape = dask.delayed(tile).image.shape[0:2]
+ di, dj = shape[0], shape[1]
  # add the Masks object for the masks corresponding to the tile
  # this assumes that the tile didn't already have any masks
  # this should work since the backend reads from image only
@@ -430,6 +450,8 @@ def generate_tiles(self, shape=3000, stride=None, pad=False, **kwargs):
  tile_slices = [slice(i, i + di), slice(j, j + dj)]
  tile.masks = self.masks.slice(tile_slices)
 
+ # TODO: end move to worker
+
  # add slide-level labels to each tile, if possible
  if self.labels is not None:
  tile.labels = self.labels

pathml/core/tile.py

@@ -7,9 +7,11 @@
 from collections import OrderedDict
 
 import anndata
+import dask
 import h5py
 import matplotlib.pyplot as plt
 import numpy as np
+from dask.delayed import Delayed
 
 import pathml.core.masks
 
@@ -21,7 +23,7 @@ class Tile:
  on labelling the top-leftmost pixel as (0, 0)
 
  Args:
- image (np.ndarray): Image array of tile
+ image (np.ndarray or dask.delayed.Delayed): Tile image or dask.delayed.Delayed object to load image
  coords (tuple): Coordinates of tile relative to the whole-slide image.
  The (i,j) coordinate system is based on labelling the top-leftmost pixel of the WSI as (0, 0).
  name (str, optional): Name of tile
@@ -60,9 +62,9 @@ def __init__(
  time_series=None,
  ):
  # check inputs
- assert isinstance(
+ assert isinstance(image, Delayed) or isinstance(
  image, np.ndarray
- ), f"image of type {type(image)} must be a np.ndarray"
+ ), f"image of type {type(image)} must be a np.ndarray or a dask.delayed.Delayed object"
  assert masks is None or isinstance(
  masks, dict
  ), f"masks is of type {type(masks)} but must be of type dict"
@@ -115,23 +117,38 @@ def __init__(
  counts, anndata.AnnData
  ), f"counts is of type {type(counts)} but must be of type anndata.AnnData or None"
 
- if masks:
- for val in masks.values():
- if val.shape[:2] != image.shape[:2]:
- raise ValueError(
- f"mask is of shape {val.shape} but must match tile shape {image.shape}"
- )
- self.masks = masks
- else:
- self.masks = OrderedDict()
-
- self.image = image
+ self._image = image
+ self.masks = masks if masks else OrderedDict()
  self.name = name
  self.coords = coords
  self.slide_type = slide_type
  self.labels = labels
  self.counts = counts
 
+ @property
+ def image(self):
+ if isinstance(self._image, Delayed):
+ image = dask.compute(self._image, scheduler="single-threaded")
+ if isinstance(image, tuple):
+ image = image[0]
+ assert isinstance(
+ image, np.ndarray
+ ), f"image of type {type(image)} must be a np.ndarray"
+ for val in self.masks.values():
+ if val.shape[:2] != image.shape[:2]:
+ raise ValueError(
+ f"mask is of shape {val.shape} but must match tile shape {image.shape}"
+ )
+ self._image = image
+ return self._image
+
+ @image.setter
+ def image(self, image):
+ assert isinstance(
+ image, np.ndarray
+ ), f"image of type {type(image)} must be a np.ndarray"
+ self._image = image
+
  def __repr__(self):
  out = []
  out.append(f"Tile(coords={self.coords}")