boundary_dataset_generate.py

# import os
# import sys
# import cv2
# import torch
import argparse
# import subprocess
import numpy as np
from glob import glob
from PIL import Image
import os.path as osp
import scipy.io as io
from tqdm import tqdm
from scipy.ndimage import distance_transform_edt, distance_transform_cdt

parser = argparse.ArgumentParser()
parser.add_argument("--datadir", dest='datadir', default='data/Inria/')
parser.add_argument("--outname", default='boundary')
# parser.add_argument('--split', nargs='+', default=['train', 'test'])
parser.add_argument('--metric', default='euc', choices=['euc', 'taxicab'])
args = parser.parse_args()

label_list = [0, 255]


def _encode_label(labelmap):
    encoded_labelmap = np.ones_like(labelmap, dtype=np.uint16) * 255
    for i, class_id in enumerate(label_list):
        encoded_labelmap[labelmap == class_id] = i

    return encoded_labelmap


def process(inp):
    """
    segfix.lib.datasets.preprocess.cityscapes.dt_offset_generator.py
    """
    (indir, outdir, basename) = inp
    # print(inp)
    labelmap = np.array(Image.open(osp.join(indir, basename)
                                   ).convert("P")).astype(np.int16)/255
    a = np.sum(labelmap == 0)
    labelmap = _encode_label(labelmap)
    labelmap = labelmap + 1
    depth_map = np.zeros(labelmap.shape, dtype=np.float32)  # (H,W,C)

    # for id in range(1, len(label_list)):
    for id in range(1, len(label_list) + 1):  # only consider outer boundary
        labelmap_i = labelmap.copy()
        labelmap_i[labelmap_i != id] = 0
        labelmap_i[labelmap_i == id] = 1
        # label: 1 background: 0
        if args.metric == 'euc':
            depth_i = distance_transform_edt(labelmap_i)
        elif args.metric == 'taxicab':
            depth_i = distance_transform_cdt(labelmap_i, metric='taxicab')
        else:
            raise RuntimeError
        depth_map += depth_i

    depth_map[depth_map > 250] = 250
    if a == (labelmap.shape[0] * labelmap.shape[1]):
        depth_map[depth_map > 0] = 250

    depth_map = depth_map.astype(np.uint8)

    # edge = depth_map.copy()
    # e1 = (depth_map > 0) & (depth_map <= 3)
    # edge[e1] = 255
    # e2 = (edge > 0) & (edge < 255)
    # edge[e2] = 0
    # ed = Image.fromarray(edge)
    # ed.save("edge.png")

    io.savemat(
        osp.join(outdir, basename.replace("tif", "mat")),
        {"depth": depth_map},
        do_compression=True,
    )


indir = osp.join(args.datadir, 'train', 'label')
outdir = osp.join(args.datadir, args.outname)
args_to_apply = [(indir, outdir, osp.basename(basename))
                 for basename in glob(osp.join(indir, "*.tif"))]
print('Processing {} files'.format(len(args_to_apply)))
for i in tqdm(range(0, len(args_to_apply))):
    process(args_to_apply[i])