123456789101112131415161718192021222324 |
- import os.path as osp
- import numpy as np
- import PIL.Image
- from labelme import logger
- from labelme.utils.draw import label_colormap
- def lblsave(filename, lbl):
- if osp.splitext(filename)[1] != '.png':
- filename += '.png'
- # Assume label ranses [-1, 254] for int32,
- # and [0, 255] for uint8 as VOC.
- if lbl.min() >= -1 and lbl.max() < 255:
- lbl_pil = PIL.Image.fromarray(lbl.astype(np.uint8), mode='P')
- colormap = label_colormap(255)
- lbl_pil.putpalette((colormap * 255).astype(np.uint8).flatten())
- lbl_pil.save(filename)
- else:
- logger.warn(
- '[%s] Cannot save the pixel-wise class label as PNG, '
- 'so please use the npy file.' % filename
- )
|