@@ -1,7 +1,7 @@
 import sys
 import unittest
 
-from tests.integration_tests import EMAIntegrationTest, LRTest, PoseEstimationDatasetIntegrationTest
+from tests.integration_tests import EMAIntegrationTest, LRTest, PoseEstimationDatasetIntegrationTest, YoloNASIntegrationTest
 
 
 class CoreIntegrationTestSuiteRunner:
@@ -19,6 +19,7 @@ class CoreIntegrationTestSuiteRunner:
         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))
+        self.integration_tests_suite.addTest(self.test_loader.loadTestsFromModule(YoloNASIntegrationTest))
 
 
 if __name__ == "__main__":

@@ -34,6 +34,7 @@ from tests.unit_tests.pose_estimation_dataset_test import TestPoseEstimationData
 from tests.unit_tests.preprocessing_unit_test import PreprocessingUnitTest
 from tests.unit_tests.quantization_utility_tests import QuantizationUtilityTest
 from tests.unit_tests.random_erase_test import RandomEraseTest
+from tests.unit_tests.replace_head_test import ReplaceHeadUnitTest
 from tests.unit_tests.strictload_enum_test import StrictLoadEnumTest
 from tests.unit_tests.train_with_intialized_param_args_test import TrainWithInitializedObjectsTest
 from tests.unit_tests.pretrained_models_unit_test import PretrainedModelsUnitTest
@@ -132,6 +133,7 @@ class CoreUnitTestSuiteRunner:
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestPoseEstimationMetrics))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestPoseEstimationDataset))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(LoadCheckpointTest))
+        self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(ReplaceHeadUnitTest))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(PreprocessingUnitTest))
 
     def _add_modules_to_end_to_end_tests_suite(self):

@@ -3,5 +3,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
+from tests.integration_tests.yolo_nas_integration_test import YoloNASIntegrationTest
 
-__all__ = ["EMAIntegrationTest", "LRTest", "PoseEstimationDatasetIntegrationTest"]
+__all__ = ["EMAIntegrationTest", "LRTest", "PoseEstimationDatasetIntegrationTest", "YoloNASIntegrationTest"]

import unittest
from super_gradients.training import models
from super_gradients.training.dataloaders import coco2017_val_yolo_nas
from super_gradients.training import Trainer
from super_gradients.training.metrics import DetectionMetrics
from super_gradients.training.models.detection_models.pp_yolo_e import PPYoloEPostPredictionCallback


class YoloNASIntegrationTest(unittest.TestCase):
    def test_yolo_nas_s_coco(self):
        trainer = Trainer("test_yolo_nas_s")
        model = models.get("yolo_nas_s", num_classes=80, pretrained_weights="coco")
        dl = coco2017_val_yolo_nas()
        metric = DetectionMetrics(
            normalize_targets=True,
            post_prediction_callback=PPYoloEPostPredictionCallback(score_threshold=0.03, nms_top_k=1000, max_predictions=300, nms_threshold=0.65),
            num_cls=80,
        )
        metric_values = trainer.test(model=model, test_loader=dl, test_metrics_list=[metric])
        self.assertAlmostEqual(metric_values[metric.map_str], 0.475, delta=0.001)

    def test_yolo_nas_m_coco(self):
        trainer = Trainer("test_yolo_nas_m")
        model = models.get("yolo_nas_m", num_classes=80, pretrained_weights="coco")
        dl = coco2017_val_yolo_nas()
        metric = DetectionMetrics(
            normalize_targets=True,
            post_prediction_callback=PPYoloEPostPredictionCallback(score_threshold=0.03, nms_top_k=1000, max_predictions=300, nms_threshold=0.65),
            num_cls=80,
        )
        metric_values = trainer.test(model=model, test_loader=dl, test_metrics_list=[metric])
        self.assertAlmostEqual(metric_values[metric.map_str], 0.5155, delta=0.001)

    def test_yolo_nas_l_coco(self):
        trainer = Trainer("test_yolo_nas_l")
        model = models.get("yolo_nas_l", num_classes=80, pretrained_weights="coco")
        dl = coco2017_val_yolo_nas()
        metric = DetectionMetrics(
            normalize_targets=True,
            post_prediction_callback=PPYoloEPostPredictionCallback(score_threshold=0.03, nms_top_k=1000, max_predictions=300, nms_threshold=0.65),
            num_cls=80,
        )
        metric_values = trainer.test(model=model, test_loader=dl, test_metrics_list=[metric])
        self.assertAlmostEqual(metric_values[metric.map_str], 0.5222, delta=0.001)


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

import os
import shutil
import unittest

import torch

import super_gradients
from super_gradients.common.object_names import Models
from super_gradients.training import models


class ReplaceHeadUnitTest(unittest.TestCase):
    def setUp(self) -> None:
        self.device = "cuda" if torch.cuda.is_available() and torch.cuda.device_count() > 0 else "cpu"
        super_gradients.init_trainer()

    def test_ppyolo_replace_head(self):
        input = torch.randn(1, 3, 640, 640).to(self.device)
        for model in [Models.PP_YOLOE_S, Models.PP_YOLOE_M, Models.PP_YOLOE_L, Models.PP_YOLOE_X]:
            model = models.get(model, pretrained_weights="coco").to(self.device).eval()
            model.replace_head(new_num_classes=100)
            (_, pred_scores), _ = model.forward(input)
            self.assertEqual(pred_scores.size(2), 100)

    def test_yolo_nas_replace_head(self):
        input = torch.randn(1, 3, 640, 640).to(self.device)
        for model in [Models.YOLO_NAS_S, Models.YOLO_NAS_M, Models.YOLO_NAS_L]:
            model = models.get(model, pretrained_weights="coco").to(self.device).eval()
            model.replace_head(new_num_classes=100)
            (_, pred_scores), _ = model.forward(input)
            self.assertEqual(pred_scores.size(2), 100)

    def tearDown(self) -> None:
        if os.path.exists("~/.cache/torch/hub/"):
            shutil.rmtree("~/.cache/torch/hub/")


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