_io.py 735 B

123456789101112131415161718192021222324
  1. import os.path as osp
  2. import numpy as np
  3. import PIL.Image
  4. from labelme import logger
  5. from labelme.utils.draw import label_colormap
  6. def lblsave(filename, lbl):
  7. if osp.splitext(filename)[1] != '.png':
  8. filename += '.png'
  9. # Assume label ranses [-1, 254] for int32,
  10. # and [0, 255] for uint8 as VOC.
  11. if lbl.min() >= -1 and lbl.max() < 255:
  12. lbl_pil = PIL.Image.fromarray(lbl.astype(np.uint8), mode='P')
  13. colormap = label_colormap(255)
  14. lbl_pil.putpalette((colormap * 255).astype(np.uint8).flatten())
  15. lbl_pil.save(filename)
  16. else:
  17. logger.warn(
  18. '[%s] Cannot save the pixel-wise class label as PNG, '
  19. 'so please use the npy file.' % filename
  20. )