Merge pull request #603 from TheChymera/master

ENH: User-specified output aspect ratio in skimage.util.montage.montage2d function

skimage/util/montage.py

@@ -6,7 +6,7 @@ from .. import exposure
 EPSILON = 1e-6
 
 
-def montage2d(arr_in, fill='mean', rescale_intensity=False):
+def montage2d(arr_in, fill='mean', rescale_intensity=False, grid_shape=None):
     """Create a 2-dimensional 'montage' from a 3-dimensional input array
     representing an ensemble of equally shaped 2-dimensional images.
 
@@ -31,13 +31,13 @@ def montage2d(arr_in, fill='mean', rescale_intensity=False):
     arr_in: ndarray, shape=[n_images, height, width]
         3-dimensional input array representing an ensemble of n_images
         of equal shape (i.e. [height, width]).
-
     fill: float or 'mean', optional
         How to fill the 2-dimensional output array when sqrt(n_images)
         is not an integer. If 'mean' is chosen, then fill = arr_in.mean().
-
     rescale_intensity: bool, optional
         Whether to rescale the intensity of each image to [0, 1].
+    grid_shape: tuple, optional
+        The desired grid shape for the montage (tiles_y, tiles_x). Tthe default aspect ratio is square.
 
     Returns
     -------
@@ -67,6 +67,13 @@ def montage2d(arr_in, fill='mean', rescale_intensity=False):
      [ 10.   11.    5.5   5.5]]
     >>> print(arr_in.mean())
     5.5
+    >>> arr_out_nonsquare = montage2d(arr_in, grid_shape=(3, 4))
+    >>> print(arr_out_nonsquare)
+    [[  0.    1.    4.    5. ]
+     [  2.    3.    6.    7. ]
+     [  8.    9.   10.   11. ]]
+    >>> print(arr_out_nonsquare.shape)
+    (3, 4)
     """
     assert arr_in.ndim == 3
 
@@ -80,19 +87,22 @@ def montage2d(arr_in, fill='mean', rescale_intensity=False):
             arr_in[i] = exposure.rescale_intensity(arr_in[i])
 
     # -- determine alpha
-    alpha = int(np.ceil(np.sqrt(n_images)))
+    if grid_shape:
+        alpha_y, alpha_x = grid_shape
+    else:
+        alpha_y = alpha_x = int(np.ceil(np.sqrt(n_images)))
 
     # -- fill missing patches
     if fill == 'mean':
         fill = arr_in.mean()
 
-    n_missing = int((alpha**2.) - n_images)
+    n_missing = int((alpha_y *  alpha_x) - n_images)
     missing = np.ones((n_missing, height, width), dtype=arr_in.dtype) * fill
     arr_out = np.vstack((arr_in, missing))
 
     # -- reshape to 2d montage, step by step
-    arr_out = arr_out.reshape(alpha, alpha, height, width)
+    arr_out = arr_out.reshape(alpha_y, alpha_x, height, width)
     arr_out = arr_out.swapaxes(1, 2)
-    arr_out = arr_out.reshape(alpha * height, alpha * width)
+    arr_out = arr_out.reshape(alpha_y * height, alpha_x * width)
 
     return arr_out

skimage/util/tests/test_montage.py

@@ -51,6 +51,23 @@ def test_shape():
 
     arr_out = montage2d(arr_in)
     assert_equal(arr_out.shape, (alpha * height, alpha * width))
+    
+    
+def test_grid_shape():
+    n_images = 6
+    height, width = 2, 2
+    arr_in = np.arange(n_images * height * width, dtype=np.float32)
+    arr_in = arr_in.reshape(n_images, height, width)
+    arr_out = montage2d(arr_in, grid_shape=(3,2))
+    correct_arr_out = np.array(
+	[[  0.,   1.,   4.,   5.],
+	 [  2.,   3.,   6.,   7.],
+	 [  8.,   9.,  12.,  13.],
+	 [ 10.,  11.,  14.,  15.],
+	 [ 16.,  17.,  20.,  21.],
+	 [ 18.,  19.,  22.,  23.]]
+    )
+    assert_array_equal(arr_out, correct_arr_out)
 
 
 def test_rescale_intensity():