How do I save label images from pydrake / Drake directly to disk?

Question

How do I write label images from Drake to disk using PIL?

If I try Image.fromarray(label), I get an error that looks like this:

  File ".../PIL/Image.py", line 2751, in fromarray
    raise TypeError("Cannot handle this data type: %s, %s" % typekey) from e
TypeError: Cannot handle this data type: (1, 1, 1), <i2

Answer 1

Some image libraries, like PIL, OpenCV, etc., may require you to reshape your input and cast to a specific dtype.

Let's say you get your label image from the following notebook:
drake/tutorials/rendering_multibody_plant.ipynb

label = sensor.label_image_output_port().Eval(sensor_context).data

Option 1.a: PIL

For PIL (at least what's available on Ubuntu 18.04), you can use np.int32.
Example:

from PIL import Image

img = label.squeeze(-1).astype(np.int32)
file = "/tmp/test.png"
Image.fromarray(img).save(file)
img_2 = np.asarray(Image.open(file))
label_2 = np.expand_dims(img_2.astype(np.int16), axis=-1)

# This should print 0, showing that you get the same image out.
print(np.max(label - label_2))

See: https://github.com/python-pillow/Pillow/issues/2970

Option 1.b: OpenCV

For OpenCV (3.4.0), you can use np.uint16, in addition to IMREAD_UNCHANGED:

import cv2

img = label.squeeze(-1).astype(np.uint16)
file = "/tmp/test.png"
cv2.imwrite(file, img)
img_2 = cv2.imread(file, cv2.IMREAD_UNCHANGED)
label_2 = np.expand_dims(img_2.astype(np.int16), axis=-1)

# This should print 0, showing that you get the same image out.
print(np.max(label - label_2))

See: http://jamesgregson.ca/16-bit-image-io-with-python.html

Example screen capture of the notebook: \

Option 2: Use *.npy or *.pkl

If you're just loading / saving for Python itself, you can just use np.save or pickle.