--- a
+++ b/slideflow/segment/data.py
@@ -0,0 +1,528 @@
+import torch
+import slideflow as sf
+import rasterio
+import numpy as np
+import shapely.affinity as sa
+
+from typing import Tuple, Union, Optional, List, Dict
+from torchvision import transforms
+from os.path import join, exists
+from rich.progress import track
+from shapely.ops import unary_union
+from shapely.geometry import Polygon
+from shapely.ops import unary_union
+from slideflow.util import path_to_name
+
+from .utils import topleft_pad_torch
+
+# -----------------------------------------------------------------------------
+
+class ThumbMaskDataset(torch.utils.data.Dataset):
+
+    def __init__(
+        self,
+        dataset: "sf.Dataset",
+        mpp: float,
+        roi_labels: List[str],
+        *,
+        mode: str = 'binary',
+    ) -> None:
+        """Dataset that generates thumbnails and ROI masks.
+
+        Args:
+            dataset (sf.Dataset): The dataset to use.
+            mpp (float): The target microns per pixel. The thumbnail will be
+                scaled to this resolution.
+            roi_labels (List[str]): The ROI labels to include in the mask.
+
+        Keyword args:
+            mode (str, optional): The mode to use for the mask. One of:
+                'binary', 'multiclass', 'multilabel'. Defaults to 'binary'.
+
+        """
+        super().__init__()
+        self.mpp = mpp
+        self.mode = mode
+        self.roi_labels = roi_labels
+
+        # Subsample dataset to only include slides with ROIs.
+        self.rois = dataset.rois()
+        slides = set(map(path_to_name, dataset.slide_paths()))
+        slides = slides.intersection(set(map(path_to_name, self.rois)))
+        dataset = dataset.filter({'slide': list(slides)})
+
+        # Prepare WSI objects (for slides with ROIs).
+        self.paths = dataset.slide_paths()
+
+    def __len__(self) -> int:
+        return len(self.paths)
+
+    def process(
+        self,
+        img: np.ndarray,
+        mask: np.ndarray
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Process the image/mask and convert to a tensor."""
+        img = torch.from_numpy(img)
+        mask = torch.from_numpy(mask)
+        return img, mask
+
+    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        # Load the image and mask.
+        path = self.paths[index]
+        wsi = sf.WSI(path, 299, 512, rois=self.rois, roi_filter_method=0.1, verbose=False)
+        output = get_thumb_and_mask(wsi, self.mpp, self.roi_labels, skip_missing=False)
+        if output is None:
+            return None
+        img = output['image']               # CHW (np.ndarray)
+        mask = output['mask'].astype(int)   # 1HW (np.ndarray)
+
+        if self.mode == 'multiclass':
+            mask = mask * np.arange(1, mask.shape[0]+1)[:, None, None]
+            mask = mask.max(axis=0)
+        elif self.mode == 'binary' and mask.ndim == 3:
+            mask = np.any(mask, axis=0)[None, :, :].astype(int)
+
+        # Process.
+        img, mask = self.process(img, mask)
+
+        return {
+            'image': img,
+            'mask': mask
+        }
+
+
+class RandomCropDataset(ThumbMaskDataset):
+
+    def __init__(self, *args, size: int = 1024, **kwargs):
+        """Dataset that generates thumbnails & ROI masks, with random crops.
+
+        Thumbnails and masks and randomly cropped and rotated together to
+        a square size of `size` pixels.
+
+        Args:
+            dataset (sf.Dataset): The dataset to use.
+            mpp (float): The target microns per pixel. The thumbnail will be
+                scaled to this resolution.
+            roi_labels (List[str]): The ROI labels to include in the mask.
+            size (int, optional): The size of the random crop. Defaults to 1024.
+
+        Keyword Args:
+            mode (str, optional): The mode to use for the mask. One of:
+                'binary', 'multiclass', 'multilabel'. Defaults to 'binary'.
+
+        """
+        super().__init__(*args, **kwargs)
+        self.size = size
+
+    def process(self, img, mask):
+        """Randomly crop/rotate the image and mask and convert to a tensor."""
+        return random_crop_and_rotate(img, mask, size=self.size)
+
+# -----------------------------------------------------------------------------
+# Buffered datasets
+
+class BufferedMaskDataset(torch.utils.data.Dataset):
+
+    def __init__(self, dataset: "sf.Dataset", source: str, *, mode: str = 'binary'):
+        """Dataset that loads buffered image and mask pairs.
+
+        Args:
+            dataset (sf.Dataset): The dataset to use.
+            source (str): The directory containing the buffered image/mask pairs.
+
+        Keyword Args:
+            mode (str, optional): The mode to use for the mask. One of:
+                'binary', 'multiclass', 'multilabel'. Defaults to 'binary'.
+
+        """
+        super().__init__()
+        if mode not in ['binary', 'multiclass', 'multilabel']:
+            raise ValueError("Invalid mode: {}. Expected one of: binary, "
+                             "multiclass, multilabel".format(mode))
+        self.dataset = dataset
+        self.mode = mode
+        self.paths = [
+            join(source, s + '.pt') for s in dataset.slides()
+            if exists(join(source, s + '.pt'))
+        ]
+
+
+    def __len__(self) -> int:
+        return len(self.paths)
+
+    def process(
+        self,
+        img: np.ndarray,
+        mask: np.ndarray
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Process the image/mask and convert to a tensor."""
+        img = torch.from_numpy(img)
+        mask = torch.from_numpy(mask)
+        return img, mask
+
+    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Load the image and mask.
+        output = torch.load(self.paths[index])
+        img = output['image']               # CHW (np.ndarray)
+        mask = output['mask'].astype(int)   # 1HW (np.ndarray)
+
+        if self.mode == 'multiclass':
+            mask = mask * np.arange(1, mask.shape[0]+1)[:, None, None]
+            mask = mask.max(axis=0)
+        elif self.mode == 'binary' and mask.ndim == 3:
+            mask = np.any(mask, axis=0)[None, :, :].astype(int)
+
+        # Process.
+        img, mask = self.process(img, mask)
+
+        return {
+            'image': img,
+            'mask': mask
+        }
+
+
+class BufferedRandomCropDataset(BufferedMaskDataset):
+
+    def __init__(self, *args, size: int = 1024, **kwargs):
+        """Dataset that loads buffered image/mask pairs and randomly crops.
+
+        Loaded thumbnails and masks and randomly cropped and rotated together to
+        a square size of `size` pixels.
+
+        Args:
+            dataset (sf.Dataset): The dataset to use.
+            source (str): The directory containing the buffered image/mask pairs.
+            size (int, optional): The size of the random crop. Defaults to 1024.
+
+        Keyword Args:
+            mode (str, optional): The mode to use for the mask. One of:
+                'binary', 'multiclass', 'multilabel'. Defaults to 'binary'.
+
+        """
+        super().__init__(*args, **kwargs)
+        self.size = size
+
+    def process(
+        self,
+        img: np.ndarray,
+        mask: np.ndarray
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Randomly crop/rotate the image and mask and convert to a tensor."""
+        return random_crop_and_rotate(img, mask, size=self.size)
+
+# -----------------------------------------------------------------------------
+
+class TileMaskDataset(torch.utils.data.Dataset):
+
+    def __init__(
+        self,
+        dataset: "sf.Dataset",
+        tile_px: int,
+        tile_um: Union[int, str],
+        *,
+        roi_labels: Optional[List[str]] = None,
+        stride_div: int = 2,
+        crop_margin: int = 0,
+        filter_method: str = 'otsu',
+        mode: str = 'binary'
+    ):
+        """Dataset that generates tiles and ROI masks from slides.
+
+        Args:
+            dataset (sf.Dataset): The dataset to use.
+            tile_px (int): The size of the tiles (in pixels).
+            tile_um (Union[int, str]): The size of the tiles (in microns).
+
+        Keyword args:
+            stride_div (int, optional): The divisor for the stride.
+                Defaults to 2.
+            crop_margin (int, optional): The number of pixels to add to the
+                tile size before random cropping to the target tile_px size.
+                Defaults to 0 (no random cropping).
+            filter_method (str, optional): The method to use for identifying
+                tiles for training. If 'roi', selects only tiles that intersect
+                with an ROI. If 'otsu', selects tiles based on an Otsu threshold
+                of the slide. Defaults to 'roi'.
+
+        """
+        super().__init__()
+
+        rois = dataset.rois()
+        slides_with_rois = [path_to_name(r) for r in rois]
+        slides = [s for s in dataset.slide_paths()
+                  if path_to_name(s) in slides_with_rois]
+        kw = dict(
+            tile_px=tile_px + crop_margin,
+            tile_um=tile_um,
+            verbose=False,
+            stride_div=stride_div
+        )
+        if roi_labels is None:
+            roi_labels = []
+        self.mode = mode
+        self.roi_labels = roi_labels
+        self.tile_px = tile_px
+        self.coords = []
+        self.all_wsi = dict()
+        self.all_wsi_with_roi = dict()
+        self.all_extract_px = dict()
+        for slide in track(slides, description="Loading slides"):
+            name = path_to_name(slide)
+            wsi = sf.WSI(slide, **kw)
+            try:
+                wsi_with_rois = sf.WSI(slide, roi_filter_method=0.1, rois=rois, **kw)
+            except Exception as e:
+                sf.log.error("Failed to load slide {}: {}".format(slide, e))
+                raise e
+
+            # Filter ROIs to only include the specified labels.
+            if self.roi_labels:
+                wsi_with_rois.rois = [roi for roi in wsi_with_rois.rois if roi.label in self.roi_labels]
+                wsi_with_rois.process_rois()
+
+            if not len(wsi_with_rois.rois):
+                continue
+            if filter_method == 'roi':
+                wsi_inner = sf.WSI(slide, roi_filter_method=0.9, rois=rois, **kw)
+                if self.roi_labels:
+                    wsi_inner.rois = [roi for roi in wsi_with_rois.rois if roi.label in self.roi_labels]
+                    wsi_inner.process_rois()
+                coords = np.argwhere(wsi_with_rois.grid & (~wsi_inner.grid)).tolist()
+            elif filter_method == 'otsu':
+                wsi.qc('otsu')
+                coords = np.argwhere(wsi.grid).tolist()
+                wsi.remove_qc()
+            elif filter_method in ['all', 'none', None]:
+                coords = np.argwhere(wsi_with_rois.grid).tolist()
+            else:
+                raise ValueError("Invalid filter method: {}. Expected one of: "
+                                 "roi, otsu".format(filter_method))
+            for c in coords:
+                self.coords.append([name] + c)
+            self.all_wsi[name] = wsi
+            self.all_wsi_with_roi[name] = wsi_with_rois
+            self.all_extract_px[name] = int(wsi.tile_um / wsi.mpp)
+
+    def __len__(self):
+        return len(self.coords)
+
+    def get_scaled_and_intersecting_polys(
+        self,
+        polys: "Polygon",
+        tile: "Polygon",
+        scale: float,
+        full_stride: int,
+        grid_idx: Tuple[int, int]
+    ):
+        """Get scaled and intersecting polygons for a given tile."""
+        gx, gy = grid_idx
+        A = polys.intersection(tile)
+
+        # Translate polygons so the intersection origin is at (0, 0)
+        B = sa.translate(A, -(full_stride*gx), -(full_stride*gy))
+
+        # Scale to the target tile size
+        C = sa.scale(B, xfact=scale, yfact=scale, origin=(0, 0))
+        return C
+
+    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Get an image and mask for a given index."""
+        slide, gx, gy = self.coords[index]
+        wsi = self.all_wsi[slide]
+        wsi_with_roi = self.all_wsi_with_roi[slide]
+        fe = self.all_extract_px[slide]
+        fs = wsi.full_stride
+        scale = wsi.tile_px / fe
+
+        # Get the image.
+        img = wsi[gx, gy].transpose(2, 0, 1)
+
+        # Get a polygon for the tile, used for determining overlapping ROIs.
+        tile = Polygon([
+            [fs*gx, fs*gy],
+            [fs*gx, (fs*gy)+fe],
+            [(fs*gx)+fe, (fs*gy)+fe],
+            [(fs*gx)+fe, fs*gy]
+        ])
+
+        # Compute the mask from ROIs.
+        if len(wsi_with_roi.rois) == 0:
+            if self.roi_labels:
+                mask = np.zeros((len(self.roi_labels), wsi.tile_px, wsi.tile_px), dtype=int)
+            else:
+                mask = np.zeros((1, wsi.tile_px, wsi.tile_px), dtype=int)
+
+        # Handle ROIs with labels (multilabel or multiclass)
+        elif self.roi_labels:
+            labeled_masks = []
+            for i, label in enumerate(self.roi_labels):
+                wsi_polys = [p.poly for p in wsi_with_roi.rois if p.label == label]
+                if len(wsi_polys) == 0:
+                    mask = np.zeros((wsi.tile_px, wsi.tile_px), dtype=int)
+                    labeled_masks.append(mask)
+                else:
+                    all_polys = unary_union(wsi_polys)
+                    polys = self.get_scaled_and_intersecting_polys(
+                        all_polys, tile, scale, fs, (gx, gy)
+                    )
+                    if isinstance(polys, Polygon) and polys.is_empty:
+                        mask = np.zeros((wsi.tile_px, wsi.tile_px), dtype=int)
+                    else:
+                        # Rasterize to an int mask.
+                        mask = rasterio.features.rasterize([polys], out_shape=[wsi.tile_px, wsi.tile_px]).astype(int)
+                    labeled_masks.append(mask)
+            mask = np.stack(labeled_masks, axis=0)
+
+        # Handle ROIs without labels (binary)
+        else:
+            # Determine the intersection at the given tile location.
+            all_polys = unary_union([p.poly for p in wsi_with_roi.rois])
+            polys = self.get_scaled_and_intersecting_polys(
+                all_polys, tile, scale, fs, (gx, gy)
+            )
+
+            if isinstance(polys, Polygon) and polys.is_empty:
+                mask = np.zeros((wsi.tile_px, wsi.tile_px), dtype=int)
+            else:
+                # Rasterize to an int mask.
+                try:
+                    mask = rasterio.features.rasterize([polys], out_shape=[wsi.tile_px, wsi.tile_px]).astype(bool).astype(np.int32)
+                except ValueError:
+                    mask = np.zeros((wsi.tile_px, wsi.tile_px), dtype=int)
+
+            # Add a dummy channel dimension.
+            mask = mask[None, :, :]
+
+        # Process according to the mode.
+        if self.mode == 'multiclass':
+            mask = mask * np.arange(1, mask.shape[0]+1)[:, None, None]
+            mask = mask.max(axis=0)
+        elif self.mode == 'binary' and mask.ndim == 3:
+            mask = np.any(mask, axis=0)[None, :, :].astype(int)
+
+        # Process.
+        img, mask = self.process(img, mask)
+
+        return {
+            'image': img,
+            'mask': mask
+        }
+
+    def process(
+        self,
+        img: np.ndarray,
+        mask: np.ndarray
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Randomly crop/rotate the image and mask and convert to a tensor."""
+        return random_crop_and_rotate(img, mask, size=self.tile_px)
+
+# -----------------------------------------------------------------------------
+
+def random_crop_and_rotate(img, mask, size):
+    if mask.ndim == 2:
+        to_squeeze = True
+        mask = mask[None, :, :]
+    else:
+        to_squeeze = False
+
+    # Convert to tensor.
+    img = torch.from_numpy(img).permute(1, 2, 0)
+    mask = torch.from_numpy(mask).permute(1, 2, 0)
+
+    # Pad to target size.
+    img = topleft_pad_torch(img, size).permute(2, 0, 1)
+    mask = topleft_pad_torch(mask, size).permute(2, 0, 1)
+
+    # Random crop.
+    i, j, h, w = transforms.RandomCrop.get_params(
+        img, output_size=(size, size))
+    img = transforms.functional.crop(img, i, j, h, w)
+    mask = transforms.functional.crop(mask, i, j, h, w)
+
+    # Random flip.
+    if np.random.rand() > 0.5:
+        img = transforms.functional.hflip(img)
+        mask = transforms.functional.hflip(mask)
+    if np.random.rand() > 0.5:
+        img = transforms.functional.vflip(img)
+        mask = transforms.functional.vflip(mask)
+
+    # Random cardinal rotation.
+    r = np.random.randint(4)
+    img = transforms.functional.rotate(img, r * 90)
+    mask = transforms.functional.rotate(mask, r * 90)
+
+    if to_squeeze:
+        mask = mask.squeeze(0)
+
+    return img, mask
+
+# -----------------------------------------------------------------------------
+
+def get_thumb_and_mask(
+    wsi: "sf.WSI",
+    mpp: float,
+    roi_labels: Optional[List[str]] = None,
+    skip_missing: bool = False
+) -> Dict[str, np.ndarray]:
+    """Get a thumbnail and segmentation mask for a slide."""
+
+    if len(wsi.rois) == 0 and skip_missing:
+        return None
+
+    # Sanity check.
+    width = int((wsi.mpp * wsi.dimensions[0]) / mpp)
+    ds = wsi.dimensions[0] / width
+    level = wsi.slide.best_level_for_downsample(ds)
+    level_dim = wsi.slide.level_dimensions[level]
+    if any([d > 10000 for d in level_dim]):
+        sf.log.warning("Large thumbnail found ({}) at level={} for {}".format(
+            level_dim, level, wsi.path
+        ))
+
+    # Get the thumbnail.
+    thumb = wsi.thumb(mpp=mpp).convert('RGB')
+    img = np.array(thumb).transpose(2, 0, 1)
+    xfact = thumb.size[1] / wsi.dimensions[1]
+    yfact = thumb.size[0] / wsi.dimensions[0]
+
+    if len(wsi.rois) == 0:
+        if roi_labels:
+            mask = np.zeros((len(roi_labels), thumb.size[1], thumb.size[0])).astype(bool)
+        else:
+            mask = np.zeros((1, thumb.size[1], thumb.size[0])).astype(bool)
+    elif roi_labels:
+        labeled_masks = []
+        for i, label in enumerate(roi_labels):
+            wsi_polys = [p.poly for p in wsi.rois if p.label == label]
+            if len(wsi_polys) == 0:
+                mask = np.zeros((thumb.size[1], thumb.size[0])).astype(bool)
+                labeled_masks.append(mask)
+            else:
+                all_polys = unary_union(wsi_polys)
+                # Scale ROIs to the thumbnail size.
+                C = sa.scale(all_polys, xfact=xfact, yfact=yfact, origin=(0, 0))
+                # Rasterize to an int mask.
+                mask = rasterio.features.rasterize([C], out_shape=(thumb.size[1], thumb.size[0])).astype(bool).astype(np.int32)
+                labeled_masks.append(mask)
+        mask = np.stack(labeled_masks, axis=0)
+
+    else:
+        all_polys = unary_union([p.poly for p in wsi.rois])
+        # Scale ROIs to the thumbnail size.
+        C = sa.scale(all_polys, xfact=xfact, yfact=yfact, origin=(0, 0))
+        # Rasterize to an int mask.
+        mask = rasterio.features.rasterize([C], out_shape=(thumb.size[1], thumb.size[0])).astype(bool)
+        # Add a dummy channel dimension.
+        mask = mask[None, :, :]
+
+    assert img.shape[1:] == mask.shape[1:], "Image and mask must have the same dimensions."
+    assert mask.ndim == 3, "Mask must have 3 dimensions (C, H, W)."
+    assert img.ndim == 3, "Image must have 3 dimensions (C, H, W)."
+
+    return {
+        'image': img,
+        'mask': mask
+    }