test_segmentor.py 641 B

123456789101112131415161718192021
  1. import unittest
  2. from segmentor import Segmentor
  3. class TestSegmentor(unittest.TestCase):
  4. def test_segmentor(self):
  5. tests = [
  6. ([1, 1, 0, 1, 0, 0, 1, 0, 1], [[0, 1, 2, 3]]),
  7. ([1, 1, 0, 1, 1, 0, 1, 0, 1], [[0, 1, 2, 3, 4, 5, 6]]),
  8. ([1, 1, 0, 1, 1, 0, 1, 1, 1], [list(range(0, 8 + 1))]),
  9. ([1, 1, 1, 0, 1], [[0, 1, 2, 3, 4]]),
  10. ([0, 0, 0, 0, 0], []),
  11. ([1] * 7 + [0] * 3, [list(range(7))])
  12. ]
  13. min_frames = 4
  14. threshold = 0.7
  15. for ex, exp in tests:
  16. self.assertEqual(exp, Segmentor._segmentor(ex, min_frames, threshold))