将标注的线段变为mask

cv2.polylines()函数的具体参数见博客

import cv2
import numpy as np
import json
import os
import glob
seq_dir = os.path.expanduser('./')
img_files=glob.glob(seq_dir + '*.json') 
print(img_files)
category_types = [ "nzm", "xgb"]
for img_file in img_files:
    imgn=str(img_file[2:7])+".jpg"
    print(str(imgn))
    img = cv2.imread(str(imgn))
    h, w = img.shape[:2]
    mask = np.zeros([h, w, 1], np.uint8)    # 创建一个大小和原图相同的空白图像

    with open(img_file, "r") as f:
        label = json.load(f)

    shapes = label["shapes"]
    for shape in shapes:
        category = shape["label"]
        points = shape["points"]
        # 填充
        points_array = np.array(points, dtype=np.int32)
        #mask = cv2.fillPoly(mask, [points_array], category_types.index(category))
        if category == "nzm":
        	# 调试时将Tom的填充颜色改为255，便于查看
            mask = cv2.polylines(mask, [points_array], False, 255, 3)
        elif category == "xgb":
            mask = cv2.polylines(mask, [points_array], False, 25, 3)
        else:
            mask = cv2.polylines(mask, [points_array], False, 0, 3)

    cv2.imwrite("mask/"+str(img_file[2:7])+'.jpg', mask)