Fix a few minor issues in image_helper

* Used ByteIO rather than StringIO is now required for PIL's Image.open
* Remove unused variable, 'alpha'
* Added support for '8UC1' and '32FC1' image types
* Only do bgr2rgb if it's definitely a BGR type

Fixes #60

Signed-off-by: Michael Jeronimo <[email protected]>

rqt_bag_plugins/src/rqt_bag_plugins/image_helper.py

@@ -32,10 +32,7 @@
 
 from __future__ import print_function
 import array
-try:
-    from cStringIO import StringIO
-except ImportError:
-    from io import StringIO
+from io import BytesIO
 import sys
 
 from PIL import Image
@@ -49,14 +46,13 @@
 def imgmsg_to_pil(img_msg, rgba=True):
     try:
         if img_msg._type == 'sensor_msgs/CompressedImage':
-            pil_img = Image.open(StringIO(img_msg.data))
-            if pil_img.mode != 'L':
+            pil_img = Image.open(BytesIO(img_msg.data))
+            if pil_img.mode.startswith('BGR'):
                 pil_img = pil_bgr2rgb(pil_img)
             pil_mode = 'RGB'
         else:
-            alpha = False
             pil_mode = 'RGB'
-            if img_msg.encoding == 'mono8':
+            if img_msg.encoding in ['mono8', '8UC1']:
                 mode = 'L'
             elif img_msg.encoding == 'rgb8':
                 mode = 'RGB'
@@ -67,21 +63,25 @@ def imgmsg_to_pil(img_msg, rgba=True):
             elif img_msg.encoding in ['bayer_rggb16', 'bayer_bggr16', 'bayer_gbrg16', 'bayer_grbg16']:
                 pil_mode = 'I;16'
                 if img_msg.is_bigendian:
-                    mode='I;16B'
+                    mode = 'I;16B'
                 else:
-                    mode='I;16L'
+                    mode = 'I;16L'
             elif img_msg.encoding == 'mono16' or img_msg.encoding == '16UC1':
                 pil_mode = 'F'
                 if img_msg.is_bigendian:
                     mode = 'F;16B'
                 else:
                     mode = 'F;16'
+            elif img_msg.encoding == '32FC1':
+                pil_mode = 'F'
+                if img_msg.is_bigendian:
+                    mode = 'F;32BF'
+                else:
+                    mode = 'F;32F'
             elif img_msg.encoding == 'rgba8':
                 mode = 'BGR'
-                alpha = True
             elif img_msg.encoding == 'bgra8':
                 mode = 'RGB'
-                alpha = True
             else:
                 raise Exception("Unsupported image format: %s" % img_msg.encoding)
             pil_img = Image.frombuffer(
@@ -90,6 +90,7 @@ def imgmsg_to_pil(img_msg, rgba=True):
         # 16 bits conversion to 8 bits
         if pil_mode == 'I;16':
             pil_img = pil_img.convert('I').point(lambda i: i * (1. / 256.)).convert('L')
+
         if pil_img.mode == 'F':
             pil_img = pil_img.point(lambda i: i * (1. / 256.)).convert('L')
             pil_img = ImageOps.autocontrast(pil_img)