integrating imgcount fix

mapbox · Sep 2, 2015 · 533628f · 533628f
1 parent 6cb1cf5
commit 533628f
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 8 deletions.
diff --git a/tests/test_tile_stitcher.py b/tests/test_tile_stitcher.py
@@ -218,7 +218,7 @@ def test_affaux(expectedAffauxs):
 def test_make_grey_imagedata():
     inputData = np.zeros((256, 256, 1), dtype=np.uint8)
 
-    imdata = stitcher.make_image_array(inputData, 256, 1)
+    imdata = stitcher.make_image_array(inputData, 256)
 
     assert imdata.shape == (4, 256, 256)
 
@@ -233,7 +233,7 @@ def test_make_grey_imagedata():
 def test_make_rgb_imagedata():
     inputData = np.zeros((256, 256, 4), dtype=np.uint8)
 
-    imdata = stitcher.make_image_array(inputData, 256, 4)
+    imdata = stitcher.make_image_array(inputData, 256)
 
     assert imdata.shape == (4, 256, 256)
     print("# OK - %s " % (inspect.stack()[0][3]))
@@ -270,6 +270,20 @@ def test_load_imagedata_grey():
     assert depth == expectedDepth
     print("# OK - %s " % (inspect.stack()[0][3]))
 
+def test_make_grey_depth2_imagedata():
+    inputData = np.zeros((256, 256), dtype=np.uint8)
+
+    imdata = stitcher.make_image_array(inputData, 256)
+
+    assert imdata.shape == (4, 256, 256)
+
+    assert np.array_equal(imdata[-1], np.zeros((256, 256), dtype=np.uint8) + 255)
+
+    assert np.array_equal(imdata[0], imdata[1])
+
+    assert np.array_equal(imdata[1], imdata[2])
+    print("# OK - %s " % (inspect.stack()[0][3]))
+
 
 def test_load_imagedata_random():
     expectedSize = int(np.random.rand() * 256)

diff --git a/tile_stitcher/__init__.py b/tile_stitcher/__init__.py
@@ -84,7 +84,11 @@ def make_window(x, y, xmin, ymin, windowsize):
 globalArgs = None
 
 
-def make_image_array(imdata, outputSize, depth):
+def make_image_array(imdata, outputSize):
+    depth = imdata.shape[-1]
+    if len(imdata.shape) == 2:
+        imdata = np.dstack([imdata])
+        depth = 1
     return np.array([
         imdata[:, :, 0 % depth],
         imdata[:, :, 1 % depth],
@@ -140,8 +144,7 @@ def streaming_tile_worker(data):
 
                     imdata = np.array(Image.open(path))
 
-                    depth = imdata.shape[-1]
-                    imdata = make_image_array(imdata, globalArgs['tileResolution'], depth)
+                    imdata = make_image_array(imdata, globalArgs['tileResolution'])
 
                     imdata = upsample(imdata, fDiff, frFaux, toFaux)
 
@@ -158,9 +161,7 @@ def streaming_tile_worker(data):
 
                 imdata = np.array(Image.open(path))
 
-                depth = imdata.shape[-1]
-
-                imdata = make_image_array(imdata, globalArgs['tileResolution'], depth)
+                imdata = make_image_array(imdata, globalArgs['tileResolution'])
 
                 window = make_window(x, y, baseX, baseY, globalArgs['tileResolution'])