import numpy as np
from numpy.testing import assert_allclose
from pybug.image import Image


def test_image_creation():
    pixels = np.ones((120, 120, 3))
    Image(pixels)


def test_2d_crop_without_mask():
    pixels = np.ones((120, 120, 3))
    im = Image(pixels)

    cropped_im, translation = im.crop(slice(10, 20), slice(50, 60))

    assert(cropped_im.image_shape == (10, 10))
    assert(cropped_im.n_channels == 3)
    assert(np.alltrue(cropped_im.image_shape))
    assert_allclose(translation.as_vector(), [0, 0, 0, 0, 10, 50])


def test_2d_crop_with_mask():
    pixels = np.ones((120, 120, 3))
    mask = np.zeros_like(pixels[..., 0])
    mask[10:100, 20:30] = 1
    im = Image(pixels, mask=mask)

    cropped_im, translation = im.crop(slice(0, 20), slice(0, 60))

    assert(cropped_im.image_shape == (20, 60))
    assert(np.alltrue(cropped_im.image_shape))

    correct_mask = np.zeros([20, 60])
    correct_mask[10:, 20:30] = 1
    assert_allclose(cropped_im.mask, correct_mask)
    assert_allclose(translation.as_vector(), [0, 0, 0, 0, 0, 0])

@@ -84,16 +84,46 @@
 t Image%0A
+from nose.tools import raises%0A
 %0A%0Adef te
@@ -1078,12 +1078,1439 @@
 , 0, 0, 0%5D)%0A
+%0A%0Adef test_3d_crop_without_mask():%0A    pixels = np.ones((120, 120, 120, 3))%0A    im = Image(pixels)%0A%0A    cropped_im, translation = im.crop(slice(10, 20), slice(50, 60),%0A                                      slice(100, 120))%0A%0A    assert(cropped_im.image_shape == (10, 10, 20))%0A    assert(cropped_im.n_channels == 3)%0A    assert(np.alltrue(cropped_im.image_shape))%0A    assert_allclose(translation.as_vector(),%0A                    %5B0, 0, 0, 0, 0, 0, 0, 0, 0, 10, 50, 100%5D)%0A%0A%0Adef test_3d_crop_with_mask():%0A    pixels = np.ones((120, 120, 120, 3))%0A    mask = np.zeros_like(pixels%5B..., 0%5D)%0A    mask%5B10:100, 20:30, 0:10%5D = 1%0A    im = Image(pixels, mask=mask)%0A%0A    cropped_im, translation = im.crop(slice(0, 20), slice(0, 60), slice(0, 10))%0A%0A    assert(cropped_im.image_shape == (20, 60, 10))%0A    assert(np.alltrue(cropped_im.image_shape))%0A%0A    correct_mask = np.zeros(%5B20, 60, 10%5D)%0A    correct_mask%5B10:, 20:30, :10%5D = 1%0A    assert_allclose(cropped_im.mask, correct_mask)%0A    assert_allclose(translation.as_vector(),%0A                    %5B0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0%5D)%0A%0A%0A@raises(ValueError)%0Adef test_4d_crop_raises_valueerror():%0A    pixels = np.ones((120, 120, 120, 120, 3))%0A    im = Image(pixels)%0A%0A    im.crop(slice(0, 20), slice(0, 60), slice(0, 10), slice(0, 10))%0A%0A%0A@raises(AssertionError)%0Adef test_crop_wrong_arg_num_raises_assertionerror():%0A    pixels = np.ones((120, 120, 3))%0A    im = Image(pixels)%0A%0A    im.crop(slice(0, 20))

Add more tests to image crop