import numpy as np
import unittest

import chainer
from chainer import testing
from chainer.testing import attr

from chainercv.links import SegNetBasic
from chainercv.utils import assert_is_semantic_segmentation_link


@testing.parameterize(
    {'train': False},
    {'train': True}
)
class TestSegNetBasic(unittest.TestCase):

    def setUp(self):
        self.n_class = 10
        self.link = SegNetBasic(n_class=self.n_class)

    def check_call(self):
        xp = self.link.xp
        x = chainer.Variable(xp.random.uniform(
            low=-1, high=1, size=(2, 3, 128, 160)).astype(np.float32))
        y = self.link(x)

        self.assertIsInstance(y, chainer.Variable)
        self.assertIsInstance(y.array, xp.ndarray)
        self.assertEqual(y.shape, (2, self.n_class, 128, 160))

    @attr.slow
    def test_call_cpu(self):
        self.check_call()

    @attr.gpu
    @attr.slow
    def test_call_gpu(self):
        self.link.to_gpu()
        self.check_call()

    @attr.slow
    def test_predict_cpu(self):
        assert_is_semantic_segmentation_link(self.link, self.n_class)

    @attr.gpu
    @attr.slow
    def test_predict_gpu(self):
        self.link.to_gpu()
        assert_is_semantic_segmentation_link(self.link, self.n_class)


class TestSegNetPretrained(unittest.TestCase):

    @attr.slow
    def test_pretrained(self):
        SegNetBasic(pretrained_model='camvid')

    @attr.slow
    def test_pretrained_n_class(self):
        SegNetBasic(n_class=11, pretrained_model='camvid')

    def test_pretrained_wrong_n_fg_class(self):
        with self.assertRaises(ValueError):
            SegNetBasic(n_class=20, pretrained_model='camvid')


testing.run_module(__name__, __file__)

import numpy as np
import unittest

import chainer
from chainer import testing
from chainer.testing import attr

from chainercv.links import SegNetBasic
from chainercv.utils import assert_is_semantic_segmentation_link


@testing.parameterize(
    {'train': False},
    {'train': True}
)
class TestSegNetBasic(unittest.TestCase):

    def setUp(self):
        self.n_class = 10
        self.link = SegNetBasic(n_class=self.n_class)

    def check_call(self):
        xp = self.link.xp
        x = chainer.Variable(xp.random.uniform(
            low=-1, high=1, size=(2, 3, 128, 160)).astype(np.float32))
        y = self.link(x)

        self.assertIsInstance(y, chainer.Variable)
        self.assertIsInstance(y.array, xp.ndarray)
        self.assertEqual(y.shape, (2, self.n_class, 128, 160))

    @attr.slow
    def test_call_cpu(self):
        self.check_call()

    @attr.gpu
    @attr.slow
    def test_call_gpu(self):
        self.link.to_gpu()
        self.check_call()

    @attr.slow
    def test_predict_cpu(self):
        assert_is_semantic_segmentation_link(self.link, self.n_class)

    @attr.gpu
    @attr.slow
    def test_predict_gpu(self):
        self.link.to_gpu()
        assert_is_semantic_segmentation_link(self.link, self.n_class)


@testing.parameterize(*testing.product({
    'n_class': [None, 5, 11],
    'pretrained_model': ['camvid'],
}))
class TestSegNetPretrained(unittest.TestCase):

    @attr.slow
    def test_pretrained(self):
        if self.pretrained_model == 'imagenet':
            valid = self.n_class is {None, 11}

        if valid:
            SegNetBasic(n_class=self.n_class,
                        pretrained_model=self.pretrained_model)
        else:
            with self.assertRaises(ValueError):
                SegNetBasic(n_class=self.n_class,
                            pretrained_model=self.pretrained_model)


testing.run_module(__name__, __file__)

improve tests of SegNetBasic