Upgrade Augmentations Pipeline in OTX + 16bit images support

application/backend/app/execution/training/otx_trainer.py

@@ -24,8 +24,8 @@
     OTXMultilabelClsDataset,
 )
 from otx.data.dataset.base import OTXDataset
+from otx.data.factory import TransformLibFactory
 from otx.data.module import OTXDataModule
-from otx.data.transform_libs.torchvision import TorchVisionTransformLib
 from otx.metrics import MetricCallable
 from otx.metrics.accuracy import MultiClassClsMetricCallable, MultiLabelClsMetricCallable
 from otx.metrics.mean_ap import MaskRLEMeanAPCallable, MeanAPCallable
@@ -229,9 +229,8 @@ def build_subset_config(subset_name: str) -> SubsetConfig:
             subset_cfg_data["input_size"] = training_config["data"]["input_size"]
             sampler_cfg_data = subset_cfg_data.pop("sampler", {})
             subset_config = SubsetConfig(sampler=SamplerConfig(**sampler_cfg_data), **subset_cfg_data)
-            subset_config.transforms = TorchVisionTransformLib.generate(  # pyrefly: ignore[bad-assignment]
-                subset_config
-            )
+            # pyrefly: ignore[missing-attribute,bad-assignment]
+            subset_config.transforms = TransformLibFactory.generate(subset_config)
             return subset_config
 
         with self._db_session_factory() as db:
@@ -291,15 +290,15 @@ def build_subset_config(subset_name: str) -> SubsetConfig:
             logger.info("Preparing {} instances for each subset", otx_dataset_class.__name__)
             otx_training_dataset = otx_dataset_class(
                 dm_subset=dm_training_dataset,
-                transforms=train_subset_config.transforms,  # pyrefly: ignore[bad-argument-type]
+                transforms=train_subset_config.transforms,  # pyrefly: ignore[missing-attribute,bad-argument-type]
             )
             otx_validation_dataset = otx_dataset_class(
                 dm_subset=dm_validation_dataset,
-                transforms=val_subset_config.transforms,  # pyrefly: ignore[bad-argument-type]
+                transforms=val_subset_config.transforms,  # pyrefly: ignore[missing-attribute,bad-argument-type]
             )
             otx_testing_dataset = otx_dataset_class(
                 dm_subset=dm_testing_dataset,
-                transforms=test_subset_config.transforms,  # pyrefly: ignore[bad-argument-type]
+                transforms=test_subset_config.transforms,  # pyrefly: ignore[missing-attribute,bad-argument-type]
             )
 
             return DatasetInfo(
@@ -368,8 +367,8 @@ def train_model(
         model_cfg["init_args"]["label_info"] = otx_datamodule.label_info.label_names
         model_cfg["init_args"]["data_input_params"] = DataInputParams(
             input_size=cast(tuple[int, int], otx_datamodule.input_size),
-            mean=otx_datamodule.input_mean,
-            std=otx_datamodule.input_std,
+            mean=otx_datamodule.input_mean if otx_datamodule.input_mean is not None else (0.0, 0.0, 0.0),
+            std=otx_datamodule.input_std if otx_datamodule.input_std is not None else (1.0, 1.0, 1.0),
         ).as_dict()
         model_parser = ArgumentParser()
         model_parser.add_subclass_arguments(OTXModel, "model", required=False, fail_untyped=False)

application/backend/tests/unit/execution/training/test_otx_trainer.py

@@ -402,31 +402,31 @@ def test_prepare_training_dataset_success(
                     "batch_size": 8,
                     "num_workers": 4,
                     "sampler": {"class_path": "torch.utils.data.RandomSampler"},
-                    "transforms": [
+                    "augmentations_cpu": [
                         {"class_path": "torchvision.transforms.v2.RandomHorizontalFlip", "init_args": {"p": 0.5}}
                     ],
                 },
                 "val_subset": {
                     "batch_size": 4,
                     "num_workers": 2,
                     "sampler": {"class_path": "torch.utils.data.RandomSampler"},
-                    "transforms": [],
+                    "augmentations_cpu": [],
                 },
                 "test_subset": {
                     "batch_size": 2,
                     "num_workers": 1,
                     "sampler": {"class_path": "torch.utils.data.RandomSampler"},
-                    "transforms": [],
+                    "augmentations_cpu": [],
                 },
             }
         }
 
-        # Mock TorchVisionTransformLib.generate to return mock transforms
+        # Mock TransformLibFactory.generate to return mock transforms
         mock_train_transforms = [Mock()]
         mock_val_transforms = [Mock()]
         mock_test_transforms = [Mock()]
 
-        with patch("app.execution.training.otx_trainer.TorchVisionTransformLib.generate") as mock_generate:
+        with patch("app.execution.training.otx_trainer.TransformLibFactory.generate") as mock_generate:
             mock_generate.side_effect = [mock_train_transforms, mock_val_transforms, mock_test_transforms]
 
             # Mock the __get_otx_dataset_class_by_task_type method to return a proper mock class
@@ -520,14 +520,17 @@ def test_prepare_training_dataset_success(
         # Verify SubsetConfig objects were created correctly
         assert dataset_info.otx_training_subset_config.batch_size == 8
         assert dataset_info.otx_training_subset_config.num_workers == 4
+        # pyrefly: ignore[missing-attribute]
         assert dataset_info.otx_training_subset_config.transforms == mock_train_transforms
 
         assert dataset_info.otx_validation_subset_config.batch_size == 4
         assert dataset_info.otx_validation_subset_config.num_workers == 2
+        # pyrefly: ignore[missing-attribute]
         assert dataset_info.otx_validation_subset_config.transforms == mock_val_transforms
 
         assert dataset_info.otx_testing_subset_config.batch_size == 2
         assert dataset_info.otx_testing_subset_config.num_workers == 1
+        # pyrefly: ignore[missing-attribute]
         assert dataset_info.otx_testing_subset_config.transforms == mock_test_transforms