add unit test for resize pre_processor
@@ -17,6 +17,7 @@ import numpy as np
import pytest
from lmnet.pre_processor import (
+ Resize,
ResizeWithJoints,
JointsToGaussianHeatmap
)
@@ -25,6 +26,25 @@ from lmnet.pre_processor import (
pytestmark = pytest.mark.usefixtures("reset_default_graph")
+def test_resize():
+ IMAGE_SIZE = [32, 32]
+ orig_image = np.zeros(shape=(1024, 512, 3), dtype=np.uint8)
+ orig_mask = np.zeros(shape=(1024, 512, 3), dtype=np.uint8)
+
+ pre_processor = Resize(IMAGE_SIZE)
+ resized = pre_processor(image=orig_image, mask=orig_mask)
+ resized_image = resized["image"]
+ resized_mask = resized["mask"]
+
+ assert isinstance(resized_image, np.ndarray)
+ assert isinstance(resized_mask, np.ndarray)
+ assert resized_image.shape[:2] == (32, 32)
+ assert resized_image.shape[2] == 3
+
+ assert resized_mask.shape[:2] == (32, 32)
+ assert resized_mask.shape[2] == 3
+
+
def test_resize_with_joints():
image = np.zeros(shape=(5, 5, 3), dtype=np.uint8)
