Dataturks生成mask图片

from skimage import draw
from skimage import io
import numpy as np
import urllib.request
import json
import logging
import os
import sys


###################  INSTALLATION NOTE #######################
##############################################################

## pip install scikit-image
## pip install numpy

###############################################################
###############################################################

#enable info logging.
logging.getLogger().setLevel(logging.INFO)

def poly2mask(blobs, c, path_to_masks_folder, h, w, label, idx):
    mask = np.zeros((h, w))
    for l in blobs:
        fill_row_coords, fill_col_coords = draw.polygon(l[1], l[0], l[2])
        mask[fill_row_coords, fill_col_coords] = 1
    io.imsave(path_to_masks_folder + "/" + str(c) + "_" + label + "_" + str(idx) + ".png", mask)


def convert_dataturks_to_masks(path_to_dataturks_annotation_json, path_to_original_images_folder, path_to_masks_folder):
    # make sure everything is setup.
    if (not os.path.isdir(path_to_original_images_folder)):
        logging.exception(
            "Please specify a valid directory path to download images, " + path_to_original_images_folder + " doesn't exist")
        return
    if (not os.path.isdir(path_to_masks_folder)):
        logging.exception(
            "Please specify a valid directory path to write mask files, " + path_to_masks_folder + " doesn't exist")
        return
    if (not os.path.exists(path_to_dataturks_annotation_json)):
        logging.exception(
            "Please specify a valid path to dataturks JSON output file, " + path_to_dataturks_annotation_json + " doesn't exist")
        return

    f = open(path_to_dataturks_annotation_json)
    train_data = f.readlines()
    train = []
    for line in train_data:
        data = json.loads(line)
        train.append(data)
    c = 0
    for objects in train:
        blobs = []
        classes = {'p':0,'pay':1,'scanning':2}
        image = objects['content'][objects['content'].rfind('_') + 1:objects['content'].rfind('.')]
        image = image.split('behavior/')[-1] + '.jpg'
        # download the images from given url
        urllib.request.urlretrieve(objects['content'], path_to_original_images_folder + "/image" + str(c) + ".jpg")
        annotations = objects['annotation']

        for annot in annotations:
            blobs = []
            label = annot['label'][0]
            print(label)
            if True:#(label != ''):
                if label not in classes:
                    classes[label] = 0

                points = annot['points']
                h = annot['imageHeight']
                w = annot['imageWidth']
                x_coord = []
                y_coord = []
                l = []
                for p in points:
                    x_coord.append(p[0] * w)
                    y_coord.append(p[1] * h)
                shape = (h, w)
                l.append(x_coord)
                l.append(y_coord)
                l.append(shape)
                blobs.append(l)
                poly2mask(blobs, image, path_to_masks_folder, annot['imageHeight'], annot['imageWidth'], label,
                          classes[label])
                classes[label] += 1
        c += 1


# if (len(sys.argv) < 4):
#     print(
#     "Please provide path to dataturks json file, path to store ground truth images and path to store mask images in this order.")
#     exit(0)
convert_dataturks_to_masks('/home/lixuan/sku_data/sku110k/yolov5_sku/weights/scan_behaviors.json','/home/lixuan/data/466yuanshi/im/imgs','/home/lixuan/data/466yuanshi/images')