# Copyright (c) Alibaba, Inc. and its affiliates.
import os
import tempfile
import unittest

import cv2
import numpy as np
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from modelscope.utils.cv.image_utils import semantic_seg_masks_to_image
from modelscope.utils.test_utils import test_level
from PIL import Image
from tests.ut_config import BASE_LOCAL_PATH


class EasyCVSegmentationPipelineTest(unittest.TestCase):
    img_path = os.path.join(BASE_LOCAL_PATH,
                            'data/test_images/image_segmentation.jpg')

    def setUp(self) -> None:
        self.task = Tasks.image_segmentation
        self.model_id = 'damo/cv_segformer-b0_image_semantic-segmentation_coco-stuff164k'

    def _internal_test_(self, model_id):
        semantic_seg = pipeline(task=Tasks.image_segmentation, model=model_id)
        outputs = semantic_seg(self.img_path)

        draw_img = semantic_seg_masks_to_image(outputs[OutputKeys.MASKS])
        with tempfile.NamedTemporaryFile(suffix='.jpg') as tmp_file:
            tmp_save_path = tmp_file.name
            cv2.imwrite(tmp_save_path, draw_img)
            print('test ' + model_id + ' DONE')

    def _internal_test_batch_(self, model_id, num_samples=2, batch_size=2):
        # TODO: support in the future
        img = np.asarray(Image.open(self.img_path))
        num_samples = num_samples
        batch_size = batch_size
        semantic_seg = pipeline(
            task=Tasks.image_segmentation,
            model=model_id,
            batch_size=batch_size)
        outputs = semantic_seg([self.img_path] * num_samples)

        self.assertEqual(semantic_seg.predict_op.batch_size, batch_size)
        self.assertEqual(len(outputs), num_samples)

        for output in outputs:
            self.assertListEqual(
                list(img.shape)[:2], list(output['seg_pred'].shape))

    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
    def test_segformer_b0(self):
        model_id = 'damo/cv_segformer-b0_image_semantic-segmentation_coco-stuff164k'
        self._internal_test_(model_id)

    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
    def test_segformer_b1(self):
        model_id = 'damo/cv_segformer-b1_image_semantic-segmentation_coco-stuff164k'
        self._internal_test_(model_id)

    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
    def test_segformer_b2(self):
        model_id = 'damo/cv_segformer-b2_image_semantic-segmentation_coco-stuff164k'
        self._internal_test_(model_id)

    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
    def test_segformer_b3(self):
        model_id = 'damo/cv_segformer-b3_image_semantic-segmentation_coco-stuff164k'
        self._internal_test_(model_id)

    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
    def test_segformer_b4(self):
        model_id = 'damo/cv_segformer-b4_image_semantic-segmentation_coco-stuff164k'
        self._internal_test_(model_id)

    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
    def test_segformer_b5(self):
        model_id = 'damo/cv_segformer-b5_image_semantic-segmentation_coco-stuff164k'
        self._internal_test_(model_id)


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