from image.sub_image import SubImageExtract
from rbm.matrix import np


def test_sub_image_rgb():
	n, c, w, h = 2, 3, 8, 8
	img = np.reshape(np.linspace(1, n*c*w*h, num=n*c*w*h, dtype=np.float64), (n, c, w, h))
	sub = SubImageExtract(4, 4, 1, 1)
	result = sub(img, 3)
	assert result.shape[1] == 3
	assert result.shape[2] == 4 and result.shape[3] == 4


def test_sub_image_gray():
	n, w, h = 2, 8, 8
	img = np.reshape(np.linspace(1, n*w*h, num=n*w*h, dtype=np.float64), (n, w, h))
	sub = SubImageExtract(4, 4, 1, 1)
	result = sub(img, 1)
	assert result.shape[1] == 1
	assert result.shape[2] == 4 and result.shape[3] == 4


if __name__ == "__main__":
	test_sub_image_rgb()
	test_sub_image_gray()
	print("Test: [passed]")