super-gradients/tests/unit_tests/preprocessing_unit_test.py

import os
import unittest
from pathlib import Path

import numpy as np
import torch

from super_gradients import Trainer
from super_gradients.common.factories.list_factory import ListFactory
from super_gradients.common.factories.processing_factory import ProcessingFactory
from super_gradients.module_interfaces import HasPreprocessingParams
from super_gradients.training import dataloaders
from super_gradients.training import models
from super_gradients.training.datasets import COCODetectionDataset
from super_gradients.training.metrics import DetectionMetrics
from super_gradients.training.models import YoloXPostPredictionCallback
from super_gradients.training.processing import (
    ReverseImageChannels,
    DetectionLongestMaxSizeRescale,
    DetectionBottomRightPadding,
    ImagePermute,
    ComposeProcessing,
)
from super_gradients.training.transforms import DetectionPaddedRescale, DetectionRGB2BGR
from super_gradients.training.utils.detection_utils import DetectionCollateFN, CrowdDetectionCollateFN


class PreprocessingUnitTest(unittest.TestCase):
    def setUp(self) -> None:
        self.mini_coco_data_dir = str(Path(__file__).parent.parent / "data" / "tinycoco")

    def test_getting_preprocessing_params(self):
        expected_image_processor = {
            "ComposeProcessing": {
                "processings": [
                    "ReverseImageChannels",
                    {"DetectionLongestMaxSizeRescale": {"output_shape": (512, 512)}},
                    {"DetectionLongestMaxSizeRescale": {"output_shape": (512, 512)}},
                    {"DetectionBottomRightPadding": {"output_shape": (512, 512), "pad_value": 114}},
                    {"ImagePermute": {"permutation": (2, 0, 1)}},
                ]
            }
        }

        train_dataset_params = {
            "data_dir": self.mini_coco_data_dir,
            "subdir": "images/train2017",
            "json_file": "instances_train2017.json",
            "cache": False,
            "input_dim": [512, 512],
            "transforms": [
                {"DetectionPaddedRescale": {"input_dim": [512, 512]}},
                {"DetectionTargetsFormatTransform": {"input_dim": [512, 512], "output_format": "LABEL_CXCYWH"}},
            ],
        }
        dataset = COCODetectionDataset(**train_dataset_params)
        preprocessing_params = dataset.get_dataset_preprocessing_params()
        self.assertEqual(len(preprocessing_params["class_names"]), 80)
        self.assertEqual(preprocessing_params["image_processor"], expected_image_processor)
        self.assertEqual(preprocessing_params["iou"], 0.65)
        self.assertEqual(preprocessing_params["conf"], 0.5)

    def test_setting_preprocessing_params_from_validation_set(self):
        train_dataset_params = {
            "data_dir": self.mini_coco_data_dir,
            "subdir": "images/train2017",
            "json_file": "instances_train2017.json",
            "cache": False,
            "input_dim": [329, 320],
            "transforms": [
                {"DetectionPaddedRescale": {"input_dim": [512, 512]}},
                {"DetectionTargetsFormatTransform": {"input_dim": [512, 512], "output_format": "LABEL_CXCYWH"}},
            ],
            "with_crowd": False,
        }

        val_dataset_params = {
            "data_dir": self.mini_coco_data_dir,
            "subdir": "images/val2017",
            "json_file": "instances_val2017.json",
            "cache": False,
            "input_dim": [329, 320],
            "transforms": [
                {"DetectionPaddedRescale": {"input_dim": [512, 512]}},
                {"DetectionTargetsFormatTransform": {"input_dim": [512, 512], "output_format": "LABEL_CXCYWH"}},
            ],
        }
        trainset = COCODetectionDataset(**train_dataset_params)
        self.assertIsInstance(trainset, HasPreprocessingParams)
        train_loader = dataloaders.get(dataset=trainset, dataloader_params={"collate_fn": DetectionCollateFN(), "num_workers": 0})

        valset = COCODetectionDataset(**val_dataset_params)
        self.assertIsInstance(valset, HasPreprocessingParams)
        valid_loader = dataloaders.get(dataset=valset, dataloader_params={"collate_fn": CrowdDetectionCollateFN(), "num_workers": 0})

        trainer = Trainer("test_setting_preprocessing_params_from_validation_set")

        detection_train_params_yolox = {
            "max_epochs": 1,
            "lr_mode": "cosine",
            "cosine_final_lr_ratio": 0.05,
            "warmup_bias_lr": 0.0,
            "warmup_momentum": 0.9,
            "initial_lr": 0.02,
            "loss": "yolox_loss",
            "criterion_params": {"strides": [8, 16, 32], "num_classes": 80},  # output strides of all yolo outputs
            "train_metrics_list": [],
            "valid_metrics_list": [DetectionMetrics(post_prediction_callback=YoloXPostPredictionCallback(), normalize_targets=True, num_cls=80)],
            "metric_to_watch": "mAP@0.50:0.95",
            "greater_metric_to_watch_is_better": True,
            "average_best_models": False,
        }

        model = models.get("yolox_s", num_classes=80)
        trainer.train(model=model, training_params=detection_train_params_yolox, train_loader=train_loader, valid_loader=valid_loader)
        processing_list = model._image_processor.processings
        self.assertTrue(isinstance(processing_list[0], ReverseImageChannels))
        self.assertTrue(isinstance(processing_list[1], DetectionLongestMaxSizeRescale))
        self.assertTrue(isinstance(processing_list[2], DetectionLongestMaxSizeRescale))
        self.assertTrue(isinstance(processing_list[3], DetectionBottomRightPadding))
        self.assertTrue(isinstance(processing_list[4], ImagePermute))
        self.assertTrue(len(processing_list), 5)
        self.assertEqual(model._default_nms_iou, 0.65)
        self.assertEqual(model._default_nms_conf, 0.5)

        checkpoint_path = os.path.join(trainer.checkpoints_dir_path, "ckpt_best.pth")
        checkpoint = torch.load(checkpoint_path, map_location="cpu")
        self.assertTrue("processing_params" in checkpoint)

    def test_setting_preprocessing_params_from_checkpoint(self):
        model = models.get("yolox_s", num_classes=80)
        self.assertTrue(model._image_processor is None)
        self.assertTrue(model._default_nms_iou is None)
        self.assertTrue(model._default_nms_conf is None)
        self.assertTrue(model._class_names is None)

        train_dataset_params = {
            "data_dir": self.mini_coco_data_dir,
            "subdir": "images/train2017",
            "json_file": "instances_train2017.json",
            "cache": False,
            "input_dim": [329, 320],
            "transforms": [
                {"DetectionPaddedRescale": {"input_dim": [512, 512]}},
                {"DetectionTargetsFormatTransform": {"input_dim": [512, 512], "output_format": "LABEL_CXCYWH"}},
            ],
            "with_crowd": False,
        }

        val_dataset_params = {
            "data_dir": self.mini_coco_data_dir,
            "subdir": "images/val2017",
            "json_file": "instances_val2017.json",
            "cache": False,
            "input_dim": [329, 320],
            "transforms": [
                {"DetectionPaddedRescale": {"input_dim": [512, 512]}},
                {"DetectionTargetsFormatTransform": {"input_dim": [512, 512], "output_format": "LABEL_CXCYWH"}},
            ],
        }
        trainset = COCODetectionDataset(**train_dataset_params)
        train_loader = dataloaders.get(dataset=trainset, dataloader_params={"collate_fn": DetectionCollateFN()})

        valset = COCODetectionDataset(**val_dataset_params)
        valid_loader = dataloaders.get(dataset=valset, dataloader_params={"collate_fn": CrowdDetectionCollateFN()})

        trainer = Trainer("save_ckpt_for")

        detection_train_params_yolox = {
            "max_epochs": 1,
            "lr_mode": "cosine",
            "cosine_final_lr_ratio": 0.05,
            "warmup_bias_lr": 0.0,
            "warmup_momentum": 0.9,
            "initial_lr": 0.02,
            "loss": "yolox_loss",
            "criterion_params": {"strides": [8, 16, 32], "num_classes": 80},  # output strides of all yolo outputs
            "train_metrics_list": [],
            "valid_metrics_list": [DetectionMetrics(post_prediction_callback=YoloXPostPredictionCallback(), normalize_targets=True, num_cls=80)],
            "metric_to_watch": "mAP@0.50:0.95",
            "greater_metric_to_watch_is_better": True,
            "average_best_models": False,
        }

        trainer.train(model=model, training_params=detection_train_params_yolox, train_loader=train_loader, valid_loader=valid_loader)

        model = models.get("yolox_s", num_classes=80, checkpoint_path=os.path.join(trainer.checkpoints_dir_path, "ckpt_best.pth"))
        processing_list = model._image_processor.processings
        self.assertTrue(isinstance(processing_list[0], ReverseImageChannels))
        self.assertTrue(isinstance(processing_list[1], DetectionLongestMaxSizeRescale))
        self.assertTrue(isinstance(processing_list[2], DetectionLongestMaxSizeRescale))
        self.assertTrue(isinstance(processing_list[3], DetectionBottomRightPadding))
        self.assertTrue(isinstance(processing_list[4], ImagePermute))
        self.assertTrue(len(processing_list), 5)
        self.assertEqual(model._default_nms_iou, 0.65)
        self.assertEqual(model._default_nms_conf, 0.5)

        checkpoint_path = os.path.join(trainer.checkpoints_dir_path, "ckpt_best.pth")
        checkpoint = torch.load(checkpoint_path, map_location="cpu")
        self.assertTrue("processing_params" in checkpoint)

    def test_processings_from_dataset_params(self):
        transforms = [DetectionRGB2BGR(prob=1), DetectionPaddedRescale(input_dim=(512, 512))]

        processings = []
        for t in transforms:
            processings += t.get_equivalent_preprocessing()

        instantiated_processing = ListFactory(ProcessingFactory()).get(processings)
        processing_pipeline = ComposeProcessing(instantiated_processing)
        result = processing_pipeline.preprocess_image(np.zeros((480, 640, 3)))
        print(result)


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