import numpy as np
import unittest

import chainer

from chainercv.experimental.links.model.fcis.utils.mask_voting \
    import mask_voting
from chainercv.utils import assert_is_bbox
from chainercv.utils import generate_random_bbox


class TestMaskVoting(unittest.TestCase):

    def setUp(self):
        n_roi = 5
        self.roi_size = 7
        self.size = (18, 24)
        self.n_class = 6
        self.bg_label = 0
        self.roi_mask_prob = np.random.uniform(
            size=(n_roi, self.roi_size, self.roi_size))
        self.roi_prob = np.random.uniform(
            size=(n_roi, self.n_class))
        self.bbox = generate_random_bbox(n_roi, self.size, 0, 18)

    def check_mask_voting(
            self, roi_mask_prob, roi_prob, bbox,
            roi_size, size, n_class, bg_label):
        xp = chainer.cuda.get_array_module(roi_mask_prob)
        roi_mask_prob, label, score, bbox = mask_voting(
            roi_mask_prob, roi_prob, bbox,
            size, n_class, 0.5, 0.3, 0.5, 0.4,
            bg_label=bg_label)

        n_roi = roi_mask_prob.shape[0]
        self.assertIsInstance(roi_mask_prob, xp.ndarray)
        self.assertEqual(roi_mask_prob.shape[1:], (roi_size, roi_size))
        self.assertTrue(
            xp.all(xp.logical_and(roi_mask_prob >= 0.0, roi_mask_prob <= 1.0)))

        self.assertIsInstance(label, xp.ndarray)
        self.assertEqual(label.shape, (n_roi, ))

        self.assertIsInstance(score, xp.ndarray)
        self.assertEqual(score.shape, (n_roi, ))

        assert_is_bbox(bbox, size)

    def test_mask_voting_cpu(self):
        self.check_mask_voting(
            self.roi_mask_prob, self.roi_prob, self.bbox,
            self.roi_size, self.size, self.n_class, self.bg_label)

add test_mask_voting.py