对于一些分割网络,需要使用到二值化的掩膜。但是制作的数据集(非公开数据集)的图像的标注是图像里物体的多边形边界。
比如:
已有的是原图和标注过的边缘(json格式的文件),想要生成下面这种二值化的掩膜。
当然,实际的情况是一张图里不止一个物体,可能有几十个。
刚开始我的想法是暴力,遍历每一个像素,之后判断该像素是否在某一个物体的边界内,如果是,就把像素变成白色,如果不是,就变成黑色。
但是这种做法十分复杂,一来要遍历每一个像素,而二来要判断该像素是否在物体边缘内(可能有好几十个物体),而且物体边缘是不规则的多边形。这样下来,光判断每像素是否在物体边界内就很麻烦了。
直到我想到了换一种思路:能否找到提取图片不规则多边形区域的方法。如果有,那么会很大的简化整个流程。
果然,网上有这个方法:使用Python/PIL裁剪多边形
代码如下:
import numpy
from PIL import Image, ImageDraw
# read image as RGB and add alpha (transparency)
im = Image.open("crop.jpg").convert("RGBA")
# convert to numpy (for convenience)
imArray = numpy.asarray(im)
# create mask
polygon = [(444,203),(623,243),(691,177),(581,26),(482,42)] # 多边形的点坐标
maskIm = Image.new('L', (imArray.shape[1], imArray.shape[0]), 0)
ImageDraw.Draw(maskIm).polygon(polygon, outline=1, fill=1)
mask = numpy.array(maskIm)
# assemble new image (uint8: 0-255)
newImArray = numpy.empty(imArray.shape,dtype='uint8')
# colors (three first columns, RGB)
newImArray[:,:,:3] = imArray[:,:,:3]
# transparency (4th column)
newImArray[:,:,3] = mask*255
# back to Image from numpy
newIm = Image.fromarray(newImArray, "RGBA")
newIm.save("out.png")
把上面边界的坐标稍微修改一下,来做个测试:
polygon = [(500,100),(500,900),(800,900),(1300,500),(1000,100)] # 多边形边界点坐标,(width,height)
其中生成掩膜的代码:
# create mask
polygon = [(500,100),(500,900),(800,900),(1300,500),(1000,100)] # 多边形边界点坐标,(width,height)
maskIm = Image.new('L', (imArray.shape[1], imArray.shape[0]), 0)
ImageDraw.Draw(maskIm).polygon(polygon, outline=1, fill=1) # outline为线条颜色,fill为填充颜色
mask = numpy.array(maskIm) # 生成了掩膜,只有多边形区域内为1,其余(含边界)全为0
既然这样可以在一个二维的0矩阵上,把一个多边形区域内置1。那么,一张图片里有多少个物体,我们就用一个for循环,把所有的多边形区域内全部置1。这样不就行了。
于是,先将代码修改为:
import numpy
from PIL import Image, ImageDraw
import cv2
# read image as RGB and add alpha (transparency)
im = Image.open("car1.jpg").convert("RGBA")
# convert to numpy (for convenience)
imArray = numpy.asarray(im)
# create mask
polygon = [(500,100),(500,900),(800,900),(1300,500),(1000,100)] # 多边形边界点坐标,(width,height)
maskIm = Image.new('L', (imArray.shape[1], imArray.shape[0]), 0)
ImageDraw.Draw(maskIm).polygon(polygon, outline=1, fill=255) # outline为线条颜色,fill为填充颜色
mask = numpy.array(maskIm) # 生成了掩膜,只有多边形区域内为1,其余(含边界)全为0
cv2.imwrite('./out2.jpg', mask)
可以看到,单独的一个掩膜已经绘制完成了。
那么我们再修改一下,使其绘制出所有的多边形。
import numpy
from PIL import Image, ImageDraw
import cv2
def draw_mask(img):
# convert to numpy (for convenience)
imArray = numpy.asarray(img)
# create mask
polygons = [[(50,10),(50,90),(80,90),(130,50),(100,10)],
[(500,100),(500,900),(800,900),(1300,500),(1000,100)] ]# 多边形边界点坐标,(width,height)
maskIm = Image.new('L', (imArray.shape[1], imArray.shape[0]), 0)
for polygon in polygons:
ImageDraw.Draw(maskIm).polygon(polygon, outline=1, fill=255) # outline为线条颜色,fill为填充颜色
mask = numpy.array(maskIm) # 生成了掩膜,只有多边形区域内为1,其余(含边界)全为0
return mask
if __name__ == '__main__':
# read image as RGB and add alpha (transparency)
im = Image.open("car1.jpg").convert("RGBA")
mask = draw_mask(im)
cv2.imwrite('out3.jpg', mask)
当然,实际的情况需要你自己读取json文件,提取出其中多边形的点坐标,存放到polygons中即可(polygons中的每一个元素是一个多边形的所有边界的的坐标)
最后,还是放一个完整的代码,具体可能还要视json文件的格式来修改:
import cv2
import os
import numpy as np
import json
from PIL import Image, ImageDraw
def draw_mask(img_name):
# 传入的是图片的名字
img = Image.open('./test/images/' + img_name + '.jpg').convert("RGBA")
imArray = np.asarray(img)
json_dir = './test/json/' + img_name + '.json'
with open(json_dir) as json_file:
json_data = json.load(json_file)
objects = json_data['shapes'] # 所有的实例
polygons = []
for object in objects:
polygon = []
for point in object['points']:
tmp = (int(point[0]), int(point[1]))
polygon.append(tmp)
polygons.append(polygon)
# polygons = [np.array(object['points'], dtype=int) for object in objects] # 把所有的object的边界的坐标存到一个list里
# polygons = np.array(polygons, dtype=int) # 这样改类型会报错,因为不是每一张图片里的实例数目都一样,导致array里各个元素维度不一样
# print(polygons[0])
maskIm = Image.new('L', (imArray.shape[1], imArray.shape[0]), 0)
for polygon in polygons:
ImageDraw.Draw(maskIm).polygon(polygon, outline=1, fill=255) # outline为线条颜色,fill为填充颜色
mask = np.array(maskIm) # 生成了掩膜,只有多边形区域内为1,其余(含边界)全为0
cv2.imwrite('./test/mask/' + img_name + '.jpg', mask)
if __name__ == '__main__':
imgs_name = [i.split('.')[0] for i in os.listdir('./test/images') if i.split('.')[-1] == 'jpg']
imgs_num = len(imgs_name)
for index, img_name in enumerate(imgs_name):
draw_mask(img_name)
print('{} / {} is done '.format(index, imgs_num))