insure that we use a NxMx4 array

Author: andrewgiessel

folium/utilities.py

@@ -300,20 +300,37 @@ def write_png(array):
         raise ValueError("Data must be NxM (mono), " + \
             "NxMx3 (rgb), or NxMx4 (rgba)")
 
+    # have to broadcast up into a full rgba array
+    array_full = np.empty((array.shape[0], array.shape[1], 4))
+    # NxM -> NxMx4
+    if array.shape[2] == 1:
+        array_full[:,:,0] = array[:,:,0]
+        array_full[:,:,1] = array[:,:,0]
+        array_full[:,:,2] = array[:,:,0]
+        array_full[:,:,3] = 1
+    # NxMx3 -> NxMx4
+    elif array.shape[2] == 3:
+        array_full[:,:,0] = array[:,:,0]
+        array_full[:,:,1] = array[:,:,1]
+        array_full[:,:,2] = array[:,:,2]
+        array_full[:,:,3] = 1
+    # NxMx4 -> keep
+    else:
+        array_full = array
+    
     # normalize to uint8 if it isn't already
-    if array.dtype != 'uint8':
-        for component in range(array.shape[2]):
-            frame = array[:,:,component]
-            array[:,:,component] = (frame / frame.max() * 255)
-        array = array.astype('uint8')
-    array = np.squeeze(array)
-    width, height = array.shape[:2]
-
-    array = array.tobytes()
+    if array_full.dtype != 'uint8':
+        for component in range(4):
+            frame = array_full[:,:,component]
+            array_full[:,:,component] = (frame / frame.max() * 255)
+        array_full = array_full.astype('uint8')
+    width, height = array_full.shape[:2]
+
+    array_full = array_full.tobytes()
 
     # reverse the vertical line order and add null bytes at the start
     width_byte_4 = width * 4
-    raw_data = b''.join(b'\x00' + array[span:span + width_byte_4]
+    raw_data = b''.join(b'\x00' + array_full[span:span + width_byte_4]
                     for span in range((height - 1) * width * 4, -1, - width_byte_4))
 
     def png_pack(png_tag, data):
@@ -327,3 +344,4 @@ def png_pack(png_tag, data):
         png_pack(b'IHDR', struct.pack("!2I5B", width, height, 8, 6, 0, 0, 0)),
         png_pack(b'IDAT', zlib.compress(raw_data, 9)),
         png_pack(b'IEND', b'')])
+