@@ -1,7 +1,7 @@
 import sys
 import unittest
 
-from tests.integration_tests import TestDataset, EMAIntegrationTest, LRTest
+from tests.integration_tests import EMAIntegrationTest, LRTest, PoseEstimationDatasetIntegrationTest
 
 
 class CoreIntegrationTestSuiteRunner:
@@ -16,9 +16,9 @@ class CoreIntegrationTestSuiteRunner:
         _add_modules_to_integration_tests_suite - Adds unit tests to the Unit Tests Test Suite
             :return:
         """
-        self.integration_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestDataset))
         self.integration_tests_suite.addTest(self.test_loader.loadTestsFromModule(EMAIntegrationTest))
         self.integration_tests_suite.addTest(self.test_loader.loadTestsFromModule(LRTest))
+        self.integration_tests_suite.addTest(self.test_loader.loadTestsFromModule(PoseEstimationDatasetIntegrationTest))
 
 
 if __name__ == "__main__":

@@ -20,6 +20,7 @@ from tests.unit_tests import (
     ResumeTrainingTest,
     CallTrainAfterTestTest,
     CrashTipTest,
+    TestTransforms,
 )
 from tests.end_to_end_tests import TestTrainer
 from tests.unit_tests.detection_utils_test import TestDetectionUtils
@@ -121,6 +122,7 @@ class CoreUnitTestSuiteRunner:
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestModelsONNXExport))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(MaxBatchesLoopBreakTest))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestTrainingUtils))
+        self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestTransforms))
 
     def _add_modules_to_end_to_end_tests_suite(self):
         """

@@ -2,5 +2,6 @@
 
 from tests.integration_tests.ema_train_integration_test import EMAIntegrationTest
 from tests.integration_tests.lr_test import LRTest
+from tests.integration_tests.pose_estimation_dataset_test import PoseEstimationDatasetIntegrationTest
 
-__all__ = ["EMAIntegrationTest", "LRTest"]
+__all__ = ["EMAIntegrationTest", "LRTest", "PoseEstimationDatasetIntegrationTest"]

import os
import unittest

import pkg_resources
from hydra import initialize_config_dir, compose
from hydra.core.global_hydra import GlobalHydra
from super_gradients.common.environment.path_utils import normalize_path
from super_gradients.training.dataloaders.dataloaders import _process_dataset_params, get_data_loader
from super_gradients.training.datasets.pose_estimation_datasets import COCOKeypointsDataset


class PoseEstimationDatasetIntegrationTest(unittest.TestCase):
    def test_datasets_instantiation(self):
        GlobalHydra.instance().clear()
        sg_recipes_dir = pkg_resources.resource_filename("super_gradients.recipes", "")
        dataset_config = os.path.join("dataset_params", "coco_pose_estimation_dekr_dataset_params")
        with initialize_config_dir(config_dir=normalize_path(sg_recipes_dir), version_base="1.2"):
            # config is relative to a module
            cfg = compose(config_name=normalize_path(dataset_config))
            train_dataset_params = _process_dataset_params(cfg, dict(), True)
            val_dataset_params = _process_dataset_params(cfg, dict(), True)

            train_dataset = COCOKeypointsDataset(**train_dataset_params)
            assert train_dataset[0] is not None

            val_dataset = COCOKeypointsDataset(**val_dataset_params)
            assert val_dataset[0] is not None

    def test_dataloaders_instantiation(self):
        train_loader = get_data_loader("coco_pose_estimation_dekr_dataset_params", COCOKeypointsDataset, train=True, dataloader_params=dict(num_workers=0))
        val_loader = get_data_loader("coco_pose_estimation_dekr_dataset_params", COCOKeypointsDataset, train=False, dataloader_params=dict(num_workers=0))

        assert next(iter(train_loader)) is not None
        assert next(iter(val_loader)) is not None


if __name__ == "__main__":
    unittest.main()

@@ -22,7 +22,7 @@ from tests.unit_tests.conv_bn_relu_test import TestConvBnRelu
 from tests.unit_tests.initialize_with_dataloaders_test import InitializeWithDataloadersTest
 from tests.unit_tests.training_params_factory_test import TrainingParamsTest
 from tests.unit_tests.config_inspector_test import ConfigInspectTest
-
+from tests.unit_tests.transforms_test import TestTransforms
 
 __all__ = [
     "CrashTipTest",
@@ -48,4 +48,5 @@ __all__ = [
     "ResumeTrainingTest",
     "CallTrainAfterTestTest",
     "ConfigInspectTest",
+    "TestTransforms",
 ]

from super_gradients.training.datasets.pose_estimation_datasets.coco_keypoints import COCOKeypointsDataset
from super_gradients.training.datasets.pose_estimation_datasets.target_generators import DEKRTargetsGenerator
from super_gradients.training.transforms.keypoint_transforms import KeypointsCompose, KeypointsRandomVerticalFlip


def test_dataset():
    target_generator = DEKRTargetsGenerator(
        output_stride=4,
        sigma=2,
        center_sigma=4,
        bg_weight=0.1,
        offset_radius=4,
    )

    dataset = COCOKeypointsDataset(
        data_dir="e:/coco2017",
        images_dir="images/train2017",
        json_file="annotations/person_keypoints_train2017.json",
        include_empty_samples=False,
        transforms=KeypointsCompose(
            [
                KeypointsRandomVerticalFlip(),
            ]
        ),
        target_generator=target_generator,
    )

    assert dataset is not None
    assert dataset[0] is not None

import unittest

import numpy as np

from super_gradients.training.transforms.keypoint_transforms import (
    KeypointsRandomHorizontalFlip,
    KeypointsRandomVerticalFlip,
    KeypointsRandomAffineTransform,
    KeypointsPadIfNeeded,
    KeypointsLongestMaxSize,
)


class TestTransforms(unittest.TestCase):
    def test_keypoints_random_affine(self):
        image = np.random.rand(640, 480, 3)
        mask = np.random.rand(640, 480)
        joints = np.random.randint(0, 480, size=(1, 17, 3))
        joints[..., 2] = 2  # all visible

        aug = KeypointsRandomAffineTransform(min_scale=0.8, max_scale=1.2, max_rotation=30, max_translate=0.5, prob=1, image_pad_value=0, mask_pad_value=0)
        aug_image, aug_mask, aug_joints = aug(image, mask, joints)

        joints_outside_image = (
            (aug_joints[:, :, 0] < 0) | (aug_joints[:, :, 1] < 0) | (aug_joints[:, :, 0] >= aug_image.shape[1]) | (aug_joints[:, :, 1] >= aug_image.shape[0])
        )
        # Ensure that keypoints outside the image are not visible
        self.assertTrue((aug_joints[joints_outside_image, 2] == 0).all())
        self.assertTrue((aug_joints[~joints_outside_image, 2] != 0).all())

    def test_keypoints_horizontal_flip(self):
        image = np.random.rand(640, 480, 3)
        mask = np.random.rand(640, 480)
        joints = np.random.randint(0, 100, size=(1, 17, 3))

        aug = KeypointsRandomHorizontalFlip(flip_index=[16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0], prob=1)
        aug_image, aug_mask, aug_joints = aug(image, mask, joints)

        np.testing.assert_array_equal(aug_image, image[:, ::-1, :])
        np.testing.assert_array_equal(aug_mask, mask[:, ::-1])
        np.testing.assert_array_equal(image.shape[1] - aug_joints[:, ::-1, 0] - 1, joints[..., 0])
        np.testing.assert_array_equal(aug_joints[:, ::-1, 1], joints[..., 1])
        np.testing.assert_array_equal(aug_joints[:, ::-1, 2], joints[..., 2])

    def test_keypoints_vertical_flip(self):
        image = np.random.rand(640, 480, 3)
        mask = np.random.rand(640, 480)
        joints = np.random.randint(0, 100, size=(1, 17, 3))

        aug = KeypointsRandomVerticalFlip(prob=1)
        aug_image, aug_mask, aug_joints = aug(image, mask, joints)

        np.testing.assert_array_equal(aug_image, image[::-1, :, :])
        np.testing.assert_array_equal(aug_mask, mask[::-1, :])
        np.testing.assert_array_equal(aug_joints[..., 0], joints[..., 0])
        np.testing.assert_array_equal(image.shape[0] - aug_joints[..., 1] - 1, joints[..., 1])
        np.testing.assert_array_equal(aug_joints[..., 2], joints[..., 2])

    def test_keypoints_pad_if_needed(self):
        image = np.random.rand(640, 480, 3)
        mask = np.random.rand(640, 480)
        joints = np.random.randint(0, 100, size=(1, 17, 3))

        aug = KeypointsPadIfNeeded(min_width=768, min_height=768, image_pad_value=0, mask_pad_value=0)
        aug_image, aug_mask, aug_joints = aug(image, mask, joints)

        self.assertEqual(aug_image.shape, (768, 768, 3))
        self.assertEqual(aug_mask.shape, (768, 768))
        np.testing.assert_array_equal(aug_joints, joints)

    def test_keypoints_longest_max_size(self):
        image = np.random.rand(640, 480, 3)
        mask = np.random.rand(640, 480)
        joints = np.random.randint(0, 480, size=(1, 17, 3))

        aug = KeypointsLongestMaxSize(max_height=512, max_width=512)
        aug_image, aug_mask, aug_joints = aug(image, mask, joints)

        self.assertEqual(aug_image.shape[:2], aug_mask.shape[:2])
        self.assertLessEqual(aug_image.shape[0], 512)
        self.assertLessEqual(aug_image.shape[1], 512)

        self.assertTrue((aug_joints[..., 0] < aug_image.shape[1]).all())
        self.assertTrue((aug_joints[..., 1] < aug_image.shape[0]).all())


if __name__ == "__main__":
    unittest.main()