WildfireDetection/datamange/createMask.py

"""
# File       : createMask.py
# Time       ：2024-04-15 9:47
# Author     ：FEANGYANG
# version    ：python 3.7
# Contact    ：1071082183@qq.com
# Description：
"""
import json
import numpy as np
import cv2
import os


def json2mask(json_path):

    # read json file
    with open(json_path, "r") as f:
        data = f.read()

    # convert str to json objs
    data = json.loads(data)

    # get the points
    points = data["shapes"][0]["points"]
    points = np.array(points, dtype=np.int32)  # tips: points location must be int32
    imh = int(data["imageHeight"])
    imw = int(data["imageWidth"])

    # create a blank image
    mask = np.zeros((imh, imw, 3), dtype=np.uint8)

    # fill the contour with 255
    cv2.fillPoly(mask, [points], (255, 255, 255))


    _mask_path = os.path.splitext(json_path)[0]
    mask_path = str(_mask_path).replace('json', 'mask')+'.jpg'


    # save the mask
    cv2.imwrite(mask_path, mask)

if __name__ == '__main__':
    path = 'data/json'
    path2 = './data/mask'
    if not os.path.exists(path2):  # 判断是否存在文件夹如果不存在则创建为文件夹
        os.makedirs(path2)

    json_paths = os.listdir(path)

    for json_path in json_paths:
        json2mask(os.path.join(path, json_path))