Skip to content
Snippets Groups Projects
test_transforms.py 1.01 KiB
Newer Older
André Anjos's avatar
André Anjos committed
# SPDX-FileCopyrightText: Copyright © 2023 Idiap Research Institute <contact@idiap.ch>
#
# SPDX-License-Identifier: GPL-3.0-or-later
"""Tests for transforms."""

import numpy
import PIL.Image
import torchvision.transforms.functional as F
André Anjos's avatar
André Anjos committed

from mednet.data.augmentations import ElasticDeformation
André Anjos's avatar
André Anjos committed


def test_elastic_deformation(datadir):
    # Get a raw sample without deformation
    data_file = str(datadir / "raw_without_elastic_deformation.png")
    raw_without_deformation = F.to_tensor(PIL.Image.open(data_file))
André Anjos's avatar
André Anjos committed

    # Elastic deforms the raw
    numpy.random.seed(seed=100)
    ed = ElasticDeformation()
André Anjos's avatar
André Anjos committed
    raw_deformed = ed(raw_without_deformation)

    # Get the same sample already deformed (with seed=100)
    data_file_2 = str(datadir / "raw_with_elastic_deformation.png")
    raw_2 = PIL.Image.open(data_file_2)

    # Compare both
    raw_deformed = (255 * numpy.asarray(raw_deformed)).astype(numpy.uint8)[
        0, :, :
    ]
André Anjos's avatar
André Anjos committed
    raw_2 = numpy.asarray(raw_2)

    numpy.testing.assert_array_equal(raw_deformed, raw_2)