diff --git a/examples/notebooks/pytorch/detcon.ipynb b/examples/notebooks/pytorch/detcon.ipynb index 527888e77..b11d92c6f 100644 --- a/examples/notebooks/pytorch/detcon.ipynb +++ b/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": [ diff --git a/examples/pytorch/detcon.py b/examples/pytorch/detcon.py index 4199bc760..7d80bf1b1 100644 --- a/examples/pytorch/detcon.py +++ b/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 ) diff --git a/lightly/transforms/detcon_transform.py b/lightly/transforms/detcon_transform.py index 82265607a..29a539a96 100644 --- a/lightly/transforms/detcon_transform.py +++ b/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]. diff --git a/tests/transforms/test_detcon_transform.py b/tests/transforms/test_detcon_transform.py index 3f33c764b..502f78923 100644 --- a/tests/transforms/test_detcon_transform.py +++ b/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)