pil_plugin.py 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259
  1. __all__ = ['imread', 'imsave']
  2. import numpy as np
  3. from PIL import Image
  4. from ...util import img_as_ubyte, img_as_uint
  5. def imread(fname, dtype=None, img_num=None, **kwargs):
  6. """Load an image from file.
  7. Parameters
  8. ----------
  9. fname : str or file
  10. File name or file-like-object.
  11. dtype : numpy dtype object or string specifier
  12. Specifies data type of array elements.
  13. img_num : int, optional
  14. Specifies which image to read in a file with multiple images
  15. (zero-indexed).
  16. kwargs : keyword pairs, optional
  17. Addition keyword arguments to pass through.
  18. Notes
  19. -----
  20. Files are read using the Python Imaging Library.
  21. See PIL docs [1]_ for a list of supported formats.
  22. References
  23. ----------
  24. .. [1] http://pillow.readthedocs.org/en/latest/handbook/image-file-formats.html
  25. """
  26. if isinstance(fname, str):
  27. with open(fname, 'rb') as f:
  28. im = Image.open(f)
  29. return pil_to_ndarray(im, dtype=dtype, img_num=img_num)
  30. else:
  31. im = Image.open(fname)
  32. return pil_to_ndarray(im, dtype=dtype, img_num=img_num)
  33. def pil_to_ndarray(image, dtype=None, img_num=None):
  34. """Import a PIL Image object to an ndarray, in memory.
  35. Parameters
  36. ----------
  37. Refer to ``imread``.
  38. """
  39. try:
  40. # this will raise an IOError if the file is not readable
  41. image.getdata()[0]
  42. except OSError as e:
  43. site = "http://pillow.readthedocs.org/en/latest/installation.html#external-libraries"
  44. pillow_error_message = str(e)
  45. error_message = (
  46. f"Could not load '{image.filename}' \n"
  47. f"Reason: '{pillow_error_message}'\n"
  48. f"Please see documentation at: {site}"
  49. )
  50. raise ValueError(error_message)
  51. frames = []
  52. grayscale = None
  53. i = 0
  54. while 1:
  55. try:
  56. image.seek(i)
  57. except EOFError:
  58. break
  59. frame = image
  60. if img_num is not None and img_num != i:
  61. image.getdata()[0]
  62. i += 1
  63. continue
  64. if image.format == 'PNG' and image.mode == 'I' and dtype is None:
  65. dtype = 'uint16'
  66. if image.mode == 'P':
  67. if grayscale is None:
  68. grayscale = _palette_is_grayscale(image)
  69. if grayscale:
  70. frame = image.convert('L')
  71. else:
  72. if image.format == 'PNG' and 'transparency' in image.info:
  73. frame = image.convert('RGBA')
  74. else:
  75. frame = image.convert('RGB')
  76. elif image.mode == '1':
  77. frame = image.convert('L')
  78. elif 'A' in image.mode:
  79. frame = image.convert('RGBA')
  80. elif image.mode == 'CMYK':
  81. frame = image.convert('RGB')
  82. if image.mode.startswith('I;16'):
  83. shape = image.size
  84. dtype = '>u2' if image.mode.endswith('B') else '<u2'
  85. if 'S' in image.mode:
  86. dtype = dtype.replace('u', 'i')
  87. frame = np.frombuffer(frame.tobytes(), dtype)
  88. frame.shape = shape[::-1]
  89. else:
  90. frame = np.array(frame, dtype=dtype)
  91. frames.append(frame)
  92. i += 1
  93. if img_num is not None:
  94. break
  95. if hasattr(image, 'fp') and image.fp:
  96. image.fp.close()
  97. if img_num is None and len(frames) > 1:
  98. return np.array(frames)
  99. elif frames:
  100. return frames[0]
  101. elif img_num:
  102. raise IndexError(f'Could not find image #{img_num}')
  103. def _palette_is_grayscale(pil_image):
  104. """Return True if PIL image in palette mode is grayscale.
  105. Parameters
  106. ----------
  107. pil_image : PIL image
  108. PIL Image that is in Palette mode.
  109. Returns
  110. -------
  111. is_grayscale : bool
  112. True if all colors in image palette are gray.
  113. """
  114. if pil_image.mode != 'P':
  115. raise ValueError('pil_image.mode must be equal to "P".')
  116. # get palette as an array with R, G, B columns
  117. # Starting in pillow 9.1 palettes may have less than 256 entries
  118. palette = np.asarray(pil_image.getpalette()).reshape((-1, 3))
  119. # Not all palette colors are used; unused colors have junk values.
  120. start, stop = pil_image.getextrema()
  121. valid_palette = palette[start : stop + 1]
  122. # Image is grayscale if channel differences (R - G and G - B)
  123. # are all zero.
  124. return np.allclose(np.diff(valid_palette), 0)
  125. def ndarray_to_pil(arr, format_str=None):
  126. """Export an ndarray to a PIL object.
  127. Parameters
  128. ----------
  129. Refer to ``imsave``.
  130. """
  131. if arr.ndim == 3:
  132. arr = img_as_ubyte(arr)
  133. mode = {3: 'RGB', 4: 'RGBA'}[arr.shape[2]]
  134. elif format_str in ['png', 'PNG']:
  135. mode = 'I;16'
  136. if arr.dtype.kind == 'f':
  137. arr = img_as_uint(arr)
  138. elif arr.max() < 256 and arr.min() >= 0:
  139. arr = arr.astype(np.uint8)
  140. mode = 'L'
  141. else:
  142. arr = img_as_uint(arr)
  143. else:
  144. arr = img_as_ubyte(arr)
  145. mode = 'L'
  146. try:
  147. array_buffer = arr.tobytes()
  148. except AttributeError:
  149. array_buffer = arr.tostring() # Numpy < 1.9
  150. if arr.ndim == 2:
  151. im = Image.new(mode, arr.T.shape)
  152. try:
  153. im.frombytes(array_buffer, 'raw', mode)
  154. except AttributeError:
  155. im.fromstring(array_buffer, 'raw', mode) # PIL 1.1.7
  156. else:
  157. image_shape = (arr.shape[1], arr.shape[0])
  158. try:
  159. im = Image.frombytes(mode, image_shape, array_buffer)
  160. except AttributeError:
  161. im = Image.fromstring(mode, image_shape, array_buffer) # PIL 1.1.7
  162. return im
  163. def imsave(fname, arr, format_str=None, **kwargs):
  164. """Save an image to disk.
  165. Parameters
  166. ----------
  167. fname : str or file-like object
  168. Name of destination file.
  169. arr : ndarray of uint8 or float
  170. Array (image) to save. Arrays of data-type uint8 should have
  171. values in [0, 255], whereas floating-point arrays must be
  172. in [0, 1].
  173. format_str : str
  174. Format to save as, this is defaulted to PNG if using a file-like
  175. object; this will be derived from the extension if fname is a string
  176. kwargs : dict
  177. Keyword arguments to the Pillow save function (or tifffile save
  178. function, for Tiff files). These are format dependent. For example,
  179. Pillow's JPEG save function supports an integer ``quality`` argument
  180. with values in [1, 95], while TIFFFile supports a ``compress``
  181. integer argument with values in [0, 9].
  182. Notes
  183. -----
  184. Use the Python Imaging Library.
  185. See PIL docs [1]_ for a list of other supported formats.
  186. All images besides single channel PNGs are converted using `img_as_uint8`.
  187. Single Channel PNGs have the following behavior:
  188. - Integer values in [0, 255] and Boolean types -> img_as_uint8
  189. - Floating point and other integers -> img_as_uint16
  190. References
  191. ----------
  192. .. [1] http://pillow.readthedocs.org/en/latest/handbook/image-file-formats.html
  193. """
  194. # default to PNG if file-like object
  195. if not isinstance(fname, str) and format_str is None:
  196. format_str = "PNG"
  197. # Check for png in filename
  198. if isinstance(fname, str) and fname.lower().endswith(".png"):
  199. format_str = "PNG"
  200. arr = np.asanyarray(arr)
  201. if arr.dtype.kind == 'b':
  202. arr = arr.astype(np.uint8)
  203. if arr.ndim not in (2, 3):
  204. raise ValueError(f"Invalid shape for image array: {arr.shape}")
  205. if arr.ndim == 3:
  206. if arr.shape[2] not in (3, 4):
  207. raise ValueError("Invalid number of channels in image array.")
  208. img = ndarray_to_pil(arr, format_str=format_str)
  209. img.save(fname, format=format_str, **kwargs)