Skip to content

Commit

Permalink
test(image): add unit tests for _to_image_tensor (#281)
Browse files Browse the repository at this point in the history
* test(image): add unit tests for _to_image_tensor

* style(image): use single-quote strings

* chore: rename variable

Co-authored-by: AlaeddineAbdessalem <alaeddine-13@live.fr>
  • Loading branch information
alessiamarcolini and alaeddine-13 authored Apr 20, 2022
1 parent 49e810e commit 6ae22db
Showing 1 changed file with 42 additions and 0 deletions.
42 changes: 42 additions & 0 deletions tests/unit/document/test_image.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
import os

import numpy as np
import pytest
from PIL import Image
import io

from docarray.document.mixins.image import _to_image_tensor


@pytest.mark.parametrize(
'width, height, output_shape',
[
(None, None, (50, 10, 3)),
(8, None, (50, 8, 3)),
(None, 30, (30, 10, 3)),
(30, 8, (8, 30, 3)),
],
)
def test_to_image_tensor_pil(rgb_image_path, width, height, output_shape):
tensor = _to_image_tensor(rgb_image_path, width, height)

assert tensor.shape == output_shape
assert isinstance(tensor, np.ndarray)


def test_to_image_tensor_blob(rgb_image_path):
with open(rgb_image_path, 'rb') as f:
blob = io.BytesIO(f.read())

tensor = _to_image_tensor(blob, None, None)

assert tensor.shape == (50, 10, 3)
assert isinstance(tensor, np.ndarray)


@pytest.fixture
def rgb_image_path(tmpdir):
img_path = os.path.join(tmpdir, 'image.png')
image = Image.new('RGB', size=(10, 50), color=(0, 0, 0))
image.save(img_path)
return img_path

0 comments on commit 6ae22db

Please sign in to comment.