controlnet_aux/tests/test_processor.py

"""Test the Processor class."""
import unittest
from PIL import Image

from controlnet_aux.processor import Processor


class TestProcessor(unittest.TestCase):
    def test_hed(self):
        processor = Processor('hed')
        image = Image.open('test_image.png')
        processed_image = processor(image)
        self.assertIsInstance(processed_image, bytes)

    def test_midas(self):
        processor = Processor('midas')
        image = Image.open('test_image.png')
        processed_image = processor(image)
        self.assertIsInstance(processed_image, bytes)

    def test_mlsd(self):
        processor = Processor('mlsd')
        image = Image.open('test_image.png')
        processed_image = processor(image)
        self.assertIsInstance(processed_image, bytes)

    def test_openpose(self):
        processor = Processor('openpose')
        image = Image.open('test_image.png')
        processed_image = processor(image)
        self.assertIsInstance(processed_image, bytes)

    def test_pidinet(self):
        processor = Processor('pidinet')
        image = Image.open('test_image.png')
        processed_image = processor(image)
        self.assertIsInstance(processed_image, bytes)

    def test_normalbae(self):
        processor = Processor('normalbae')
        image = Image.open('test_image.png')
        processed_image = processor(image)
        self.assertIsInstance(processed_image, bytes)

    def test_lineart(self):
        processor = Processor('lineart')
        image = Image.open('test_image.png')
        processed_image = processor(image)
        self.assertIsInstance(processed_image, bytes)

    def test_lineart_coarse(self):
        processor = Processor('lineart_coarse')
        image = Image.open('test_image.png')
        processed_image = processor(image)
        self.assertIsInstance(processed_image, bytes)

    def test_lineart_anime(self):
        processor = Processor('lineart_anime')
        image = Image.open('test_image.png')
        processed_image = processor(image)
        self.assertIsInstance(processed_image, bytes)

    def test_canny(self):
        processor = Processor('canny')
        image = Image.open('test_image.png')
        processed_image = processor(image)
        self.assertIsInstance(processed_image, bytes)

    def test_content_shuffle(self):
        processor = Processor('content_shuffle')
        image = Image.open('test_image.png')
        processed_image = processor(image)
        self.assertIsInstance(processed_image, bytes)

    def test_zoe(self):
        processor = Processor('zoe')
        image = Image.open('test_image.png')
        processed_image = processor(image)
        self.assertIsInstance(processed_image, bytes)

    def test_mediapipe_face(self):
        processor = Processor('mediapipe_face')
        image = Image.open('test_image.png')
        processed_image = processor(image)
        self.assertIsInstance(processed_image, bytes)


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