Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,782 Bytes
29d49a9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import io
import numpy as np
import pytest
from PIL import Image
from controlnet_aux.processor import MODELS, Processor
@pytest.fixture(params=[
'scribble_hed',
'softedge_hed',
'scribble_hedsafe',
'softedge_hedsafe',
'depth_midas',
'mlsd',
'openpose',
'openpose_hand',
'openpose_face',
'openpose_faceonly',
'openpose_full',
'scribble_pidinet',
'softedge_pidinet',
'scribble_pidsafe',
'softedge_pidsafe',
'normal_bae',
'lineart_coarse',
'lineart_realistic',
'lineart_anime',
'canny',
'shuffle',
'depth_zoe',
'depth_leres',
'depth_leres++',
'mediapipe_face'
])
def processor(request):
return Processor(request.param)
def test_processor_init(processor):
assert isinstance(processor.processor, MODELS[processor.processor_id]['class'])
assert isinstance(processor.params, dict)
def test_processor_call(processor):
# Load test image
with open('test_image.png', 'rb') as f:
image_bytes = f.read()
image = Image.open(io.BytesIO(image_bytes))
# Output size
resolution = 512
W, H = image.size
H = float(H)
W = float(W)
k = float(resolution) / min(H, W)
H *= k
W *= k
H = int(np.round(H / 64.0)) * 64
W = int(np.round(W / 64.0)) * 64
# Test processing
processed_image = processor(image)
assert isinstance(processed_image, Image.Image)
assert processed_image.size == (W, H)
def test_processor_call_bytes(processor):
# Load test image
with open('test_image.png', 'rb') as f:
image_bytes = f.read()
# Test processing
processed_image_bytes = processor(image_bytes, to_pil=False)
assert isinstance(processed_image_bytes, bytes)
assert len(processed_image_bytes) > 0 |