Python实现RLE格式与PNG格式互转

介绍

在机器视觉领域的深度学习中，每个数据集都有一份标注好的数据用于训练神经网络。

为了节省空间，很多数据集的标注文件使用RLE的格式。

但是神经网络的输入一定是一张图片，为此必须把RLE格式的文件转变为图像格式。

图像格式主要又分为 .jpg 和 .png 两种格式，其中label数据一定不能使用 .jpg，因为它因为压缩算算法的原因，会造成图像失真，图像各个像素的值可能会发生变化。分割任务的数据集的 label 图像中每一个像素都代表了该像素点所属的类别，所以这样的失真是无法接受的。为此只能使用 .png 格式作为label，pascol voc 和 coco 数据集正是这样做的。

1.PNG2RLE

PNG格式转RLE格式

#!---- coding: utf- ---- import numpy as np

def rle_encode(binary_mask):
   '''
   binary_mask: numpy array, 1 - mask, 0 - background
   Returns run length as string formated
   '''
   pixels = binary_mask.flatten()
   pixels = np.concatenate([[0], pixels, [0]])
   runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
   runs[1::2] -= runs[::2]
   return ' '.join(str(x) for x in runs)

2.RLE2PNG

RLE格式转PNG格式

#!--*-- coding: utf- --*--
import numpy as np

def rle_decode(mask_rle, shape):
   '''
   mask_rle: run-length as string formated (start length)
   shape: (height,width) of array to return
   Returns numpy array, 1 - mask, 0 - background
   '''
   s = mask_rle.split()
   starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]
   starts -= 1
   ends = starts + lengths
   binary_mask = np.zeros(shape[0] * shape[1], dtype=np.uint8)
   for lo, hi in zip(starts, ends):
       binary_mask[lo:hi] = 1
   return binary_mask.reshape(shape)

3.示例

'''
RLE: Run-Length Encode
'''
from PIL import Image
import numpy as np

def __main__():
   maskfile = '/path/to/test.png'
   mask = np.array(Image.open(maskfile))
   binary_mask = mask.copy()
   binary_mask[binary_mask <= 127] = 0
   binary_mask[binary_mask > 127] = 1

# encode
   rle_mask = rle_encode(binary_mask)

# decode
   binary_mask_decode = self.rle_decode(rle_mask, binary_mask.shape[:2])

4.完整代码如下

'''
RLE: Run-Length Encode
'''
#!--*-- coding: utf- --*--
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

# M1:
class general_rle(object):
   '''
   ref.: https://www.kaggle.com/stainsby/fast-tested-rle
   '''
   def __init__(self):
       pass

def rle_encode(self, binary_mask):
       pixels = binary_mask.flatten()
       # We avoid issues with '1' at the start or end (at the corners of
       # the original image) by setting those pixels to '0' explicitly.
       # We do not expect these to be non-zero for an accurate mask,
       # so this should not harm the score.
       pixels[0] = 0
       pixels[-1] = 0
       runs = np.where(pixels[1:] != pixels[:-1])[0] + 2
       runs[1::2] = runs[1::2] - runs[:-1:2]
       return runs

def rle_to_string(self, runs):
       return ' '.join(str(x) for x in runs)

def check(self):
       test_mask = np.asarray([[0, 0, 0, 0],
                               [0, 0, 1, 1],
                               [0, 0, 1, 1],
                               [0, 0, 0, 0]])
       assert rle_to_string(rle_encode(test_mask)) == '7 2 11 2'

# M2:
class binary_mask_rle(object):
   '''
   ref.: https://www.kaggle.com/paulorzp/run-length-encode-and-decode
   '''
   def __init__(self):
       pass

def rle_encode(self, binary_mask):
       '''
       binary_mask: numpy array, 1 - mask, 0 - background
       Returns run length as string formated
       '''
       pixels = binary_mask.flatten()
       pixels = np.concatenate([[0], pixels, [0]])
       runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
       runs[1::2] -= runs[::2]
       return ' '.join(str(x) for x in runs)

def rle_decode(self, mask_rle, shape):
       '''
       mask_rle: run-length as string formated (start length)
       shape: (height,width) of array to return
       Returns numpy array, 1 - mask, 0 - background
       '''
       s = mask_rle.split()
       starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]
       starts -= 1
       ends = starts + lengths
       binary_mask = np.zeros(shape[0] * shape[1], dtype=np.uint8)
       for lo, hi in zip(starts, ends):
           binary_mask[lo:hi] = 1
       return binary_mask.reshape(shape)

def check(self):
       maskfile = '/path/to/test.png'
       mask = np.array(Image.open(maskfile))
       binary_mask = mask.copy()
       binary_mask[binary_mask <= 127] = 0
       binary_mask[binary_mask > 127] = 1

# encode
       rle_mask = self.rle_encode(binary_mask)

# decode
       binary_mask2 = self.rle_decode(rle_mask, binary_mask.shape[:2])

来源：https://juejin.cn/post/7132869631328403493