feat(transforms): add mask_fn to DetConSTransform for callable segmentation

examples/notebooks/pytorch/detcon.ipynb

@@ -43,9 +43,7 @@
     "import torch\n",
     "import torch.nn.functional as F\n",
     "import torchvision\n",
-    "from torch import nn\n",
-    "from torchvision.transforms.v2 import ToImage\n",
-    "from torchvision.tv_tensors import Mask"
+    "from torch import nn"
    ]
   },
   {
@@ -122,11 +120,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "num_cls = 25\n",
-    "if SCIKIT_IMAGE_INSTALLED:\n",
-    "    _detcons_transform = DetConSTransform(input_size=96)\n",
-    "else:\n",
-    "    _detcons_transform = DetConSTransform(grid_size=(5, 5), input_size=96)"
+    "num_cls = 25"
    ]
   },
   {
@@ -136,7 +130,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model = DetConS(backbone, num_cls=num_cls)"
+    "if SCIKIT_IMAGE_INSTALLED:\n",
+    "\n",
+    "    def felzenszwalb_mask(image):\n",
+    "        # Return the integer labels as a tensor; DetConSTransform wraps them into a mask.\n",
+    "        segments = felzenszwalb(np.asarray(image), scale=100, sigma=0.5, min_size=20)\n",
+    "        segments = np.clip(segments, 0, num_cls - 1)\n",
+    "        return torch.from_numpy(segments.astype(np.int64))\n",
+    "\n",
+    "    transform = DetConSTransform(mask_fn=felzenszwalb_mask, input_size=96)\n",
+    "else:\n",
+    "    transform = DetConSTransform(grid_size=(5, 5), input_size=96)"
    ]
   },
   {
@@ -146,8 +150,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
-    "model.to(device)"
+    "model = DetConS(backbone, num_cls=num_cls)"
    ]
   },
   {
@@ -157,7 +160,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "_to_image = ToImage()"
+    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "model.to(device)"
    ]
   },
   {
@@ -166,26 +170,6 @@
    "id": "12",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "def transform(pil_img):\n",
-    "    tv_img = _to_image(pil_img)\n",
-    "    if SCIKIT_IMAGE_INSTALLED:\n",
-    "        segments = felzenszwalb(\n",
-    "            np.array(pil_img), scale=100, sigma=0.5, min_size=20\n",
-    "        ).astype(np.int64)\n",
-    "        segments = np.clip(segments, 0, num_cls - 1)\n",
-    "        mask = Mask(torch.from_numpy(segments).unsqueeze(0))\n",
-    "    else:\n",
-    "        mask = Mask(torch.zeros(1, *tv_img.shape[-2:], dtype=torch.int64))\n",
-    "    return _detcons_transform(tv_img, mask)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "13",
-   "metadata": {},
-   "outputs": [],
    "source": [
     "dataset = torchvision.datasets.CIFAR10(\n",
     "    \"datasets/cifar10\", download=True, transform=transform\n",
@@ -197,7 +181,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "14",
+   "id": "13",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -213,7 +197,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "15",
+   "id": "14",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -224,7 +208,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "16",
+   "id": "15",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -235,7 +219,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "17",
+   "id": "16",
    "metadata": {},
    "outputs": [],
    "source": [

examples/pytorch/detcon.py

@@ -13,8 +13,6 @@
 import torch.nn.functional as F
 import torchvision
 from torch import nn
-from torchvision.transforms.v2 import ToImage
-from torchvision.tv_tensors import Mask
 
 from lightly.loss import DetConSLoss
 from lightly.models import utils
@@ -54,32 +52,24 @@ def forward(self, x, mask):
 backbone = nn.Sequential(*list(resnet.children())[:-2])
 
 num_cls = 25
+
 if SCIKIT_IMAGE_INSTALLED:
-    _detcons_transform = DetConSTransform(input_size=96)
+
+    def felzenszwalb_mask(image):
+        # Return the integer labels as a tensor; DetConSTransform wraps them into a mask.
+        segments = felzenszwalb(np.asarray(image), scale=100, sigma=0.5, min_size=20)
+        segments = np.clip(segments, 0, num_cls - 1)
+        return torch.from_numpy(segments.astype(np.int64))
+
+    transform = DetConSTransform(mask_fn=felzenszwalb_mask, input_size=96)
 else:
-    _detcons_transform = DetConSTransform(grid_size=(5, 5), input_size=96)
+    transform = DetConSTransform(grid_size=(5, 5), input_size=96)
 
 model = DetConS(backbone, num_cls=num_cls)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
-_to_image = ToImage()
-
-
-def transform(pil_img):
-    tv_img = _to_image(pil_img)
-    if SCIKIT_IMAGE_INSTALLED:
-        segments = felzenszwalb(
-            np.array(pil_img), scale=100, sigma=0.5, min_size=20
-        ).astype(np.int64)
-        segments = np.clip(segments, 0, num_cls - 1)
-        mask = Mask(torch.from_numpy(segments).unsqueeze(0))
-    else:
-        mask = Mask(torch.zeros(1, *tv_img.shape[-2:], dtype=torch.int64))
-    return _detcons_transform(tv_img, mask)
-
-
 dataset = torchvision.datasets.CIFAR10(
     "datasets/cifar10", download=True, transform=transform
 )

lightly/transforms/detcon_transform.py

@@ -1,4 +1,8 @@
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+from PIL.Image import Image
+from torch import Tensor
+from torchvision.tv_tensors import Mask
 
 from lightly.transforms.add_grid_transform import AddGridTransform
 from lightly.transforms.multi_view_transform_v2 import MultiViewTransformV2
@@ -28,8 +32,16 @@ class DetConSTransform(MultiViewTransformV2):
         - RandomGrayscale
         - GaussianBlur (only for the first view)
 
-    Can additionally apply a segmentation of the image into a regular grid if not provided
-    with a pre-segmented image.
+    The segmentation mask can be provided in one of three ways:
+        - Passing ``grid_size`` segments the image into a regular grid.
+        - Passing ``mask_fn`` derives the mask from each image on the fly, for example
+          with an unsupervised segmentation algorithm such as
+          ``skimage.segmentation.felzenszwalb``. The mask is generated once per image
+          and shared by both views, so their region correspondence is preserved.
+        - Passing neither expects the caller to supply a pre-segmented mask alongside
+          the image.
+
+    ``grid_size`` and ``mask_fn`` are mutually exclusive.
 
     References:
         - [0] DetCon, 2021, https://arxiv.org/abs/2103.10957
@@ -38,7 +50,15 @@ class DetConSTransform(MultiViewTransformV2):
 
     Attributes:
         grid_size: Size of the grid segmentation as a tuple (num_rows, num_cols), or None
-            if the segmentation mask is to be provided by the user.
+            if the segmentation mask is provided by the user or by ``mask_fn``.
+        mask_fn: Optional callable that maps an image to a segmentation mask. It receives
+            the image passed to the transform and returns the integer segmentation labels
+            as a ``Tensor`` or PIL image, for example the output of an unsupervised
+            segmenter such as ``skimage.segmentation.felzenszwalb`` wrapped as a tensor.
+            The transform wraps the result into a ``torchvision.tv_tensors.Mask``
+            internally. It must be picklable to run in ``DataLoader`` worker processes,
+            so use a module-level function or ``functools.partial`` rather than a lambda.
+            Mutually exclusive with ``grid_size``.
         gaussian_blur_t1:
             Probability of applying Gaussian blur to the first view.
         gaussian_blur_t2:
@@ -84,6 +104,9 @@ class DetConSTransform(MultiViewTransformV2):
     def __init__(
         self,
         grid_size: Optional[Tuple[int, int]] = None,
+        mask_fn: Optional[
+            Callable[[Union[Image, Tensor]], Union[Image, Tensor]]
+        ] = None,
         gaussian_blur_t1: float = 1.0,
         gaussian_blur_t2: float = 0.0,
         input_size: Union[Tuple[int, int], int] = 224,
@@ -103,7 +126,12 @@ def __init__(
         rr_degrees: Union[float, Tuple[float, float]] = 0.0,
         normalize: Union[None, Dict[str, List[float]]] = IMAGENET_NORMALIZE,
     ) -> None:
+        if grid_size is not None and mask_fn is not None:
+            raise ValueError(
+                "`grid_size` and `mask_fn` are mutually exclusive; provide at most one."
+            )
         self.grid_size = grid_size
+        self.mask_fn = mask_fn
 
         tr1: List[Union[AddGridTransform, DetConSViewTransform]] = []
         tr2: List[Union[AddGridTransform, DetConSViewTransform]] = []
@@ -163,6 +191,33 @@ def __init__(
 
         super().__init__(transforms=[T.Compose(tr1), T.Compose(tr2)])
 
+    def __call__(self, *args: Any) -> List[Any]:
+        """Creates two views of the input, generating the mask via ``mask_fn`` if set.
+
+        When ``mask_fn`` is provided, the transform is called with the image only. The
+        mask is generated once from that image and wrapped into a
+        ``torchvision.tv_tensors.Mask``, then passed through both view pipelines
+        together with the image, so both views share the same segmentation. Otherwise
+        the arguments are forwarded unchanged (image plus an optional pre-segmented or
+        grid-placeholder mask).
+
+        Args:
+            *args: Either a single image (when ``mask_fn`` is set) or an image together
+                with a mask, compatible with torchvision transforms v2.
+
+        Returns:
+            A list of two views, where each view is a transformed version of the input.
+        """
+        if self.mask_fn is not None:
+            image = args[0]
+            mask = Mask(self.mask_fn(image))
+            # Segmenters such as skimage.felzenszwalb return (H, W) labels; add a
+            # channel dimension so the mask matches the (1, H, W) grid-path convention.
+            if mask.ndim == 2:
+                mask = Mask(mask.unsqueeze(0))
+            return super().__call__(image, mask)
+        return super().__call__(*args)
+
 
 class DetConSViewTransform:
     """Transforms an image into a view for DetConS [0, 1].

tests/transforms/test_detcon_transform.py

@@ -59,3 +59,31 @@ def test_generate_grid_mask(self, img: Image, mask: Mask) -> None:
         assert mask_tr2.shape == (1, 256, 256)
 
         assert (mask_tr1.unique() == torch.arange(4 * 4)).all()
+
+    def test_mask_fn(self, img: Image) -> None:
+        # mask_fn receives the image and returns raw integer labels, mirroring an
+        # unsupervised segmenter such as skimage.segmentation.felzenszwalb.
+        def segment(image: Image) -> torch.Tensor:
+            h, w = image.shape[-2:]
+            labels = torch.zeros((h, w), dtype=torch.int64)
+            labels[:, w // 2 :] = 1
+            return labels
+
+        # deactivate anything that could change the mask
+        tr = DetConSTransform(
+            mask_fn=segment, input_size=(256, 256), min_scale=1.0, rr_prob=0.0
+        )
+
+        # the transform is called with the image only; the mask is generated internally
+        (img_tr1, mask_tr1), (img_tr2, mask_tr2) = tr(img)
+
+        assert img_tr1.shape == (3, 256, 256)
+        assert mask_tr1.shape == (1, 256, 256)
+        assert img_tr2.shape == (3, 256, 256)
+        assert mask_tr2.shape == (1, 256, 256)
+
+        assert (mask_tr1.unique() == torch.tensor([0, 1])).all()
+
+    def test_grid_size_and_mask_fn_are_mutually_exclusive(self) -> None:
+        with pytest.raises(ValueError):
+            DetConSTransform(grid_size=(4, 4), mask_fn=lambda image: image)