# ================================================================
# MIT License
# Copyright (c) 2021 edwardyehuang (https://github.com/edwardyehuang)
# ================================================================

import os, sys

rootpath = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, os.pardir))

sys.path.insert(1, rootpath)

import tensorflow as tf

import numpy as np
from PIL import Image

from absl import app

from absl import flags
from common_flags import FLAGS

from ids.voc2012 import get_colormap as get_voc2012_colormap
from ids.cityscapes_fine import get_colormap as get_cityscapes_colormap

flags.DEFINE_string("input_dir", None, "input dir path")
flags.DEFINE_string("output_dir", None, "output dir path")
flags.DEFINE_string("colormap", "voc2012", "colormap name")
flags.DEFINE_integer("ignore_label", 255, "ignore label")


def apply_colormap_to_dir(input_dir, output_dir=None, colormap=None):

    colormap = colormap.astype(np.uint8)

    counter = 0

    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    for filename in tf.io.gfile.listdir(input_dir):

        input_path = os.path.join(input_dir, filename)
        output_path = os.path.join(output_dir, filename)

        img = Image.open(input_path)

        if img.mode != "L" and img.mode != "P":
            continue

        img = img.convert("P")

        img.putpalette(colormap)
        img.save(output_path, format="PNG")

        counter += 1

        tf.print("Processed {}".format(counter))


def main(argv):

    colormap_name = FLAGS.colormap

    colormap_name = colormap_name.lower()

    if colormap_name == "voc2012":
        colormap = get_voc2012_colormap()
    elif colormap_name == "cityscapes":
        colormap = get_cityscapes_colormap()
    else:
        raise ValueError(f"Not support colormap = {colormap_name}")

    if FLAGS.ignore_label == 0:
        colormap = colormap[1:]

    apply_colormap_to_dir(FLAGS.input_dir, FLAGS.output_dir, colormap=colormap)


if __name__ == "__main__":

    app.run(main)

[((10, 0, 10, 28), 'sys.path.insert', 'sys.path.insert', ({(10, 16, 10, 17): '(1)', (10, 19, 10, 27): 'rootpath'}, {}), '(1, rootpath)', False, 'import os, sys\n'), ((25, 0, 25, 56), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', ({(25, 20, 25, 31): '"""input_dir"""', (25, 33, 25, 37): 'None', (25, 39, 25, 55): '"""input dir path"""'}, {}), "('input_dir', None, 'input dir path')", False, 'from absl import flags\n'), ((26, 0, 26, 58), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', ({(26, 20, 26, 32): '"""output_dir"""', (26, 34, 26, 38): 'None', (26, 40, 26, 57): '"""output dir path"""'}, {}), "('output_dir', None, 'output dir path')", False, 'from absl import flags\n'), ((27, 0, 27, 59), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', ({(27, 20, 27, 30): '"""colormap"""', (27, 32, 27, 41): '"""voc2012"""', (27, 43, 27, 58): '"""colormap name"""'}, {}), "('colormap', 'voc2012', 'colormap name')", False, 'from absl import flags\n'), ((28, 0, 28, 57), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(28, 21, 28, 35): '"""ignore_label"""', (28, 37, 28, 40): '(255)', (28, 42, 28, 56): '"""ignore label"""'}, {}), "('ignore_label', 255, 'ignore label')", False, 'from absl import flags\n'), ((40, 20, 40, 50), 'tensorflow.io.gfile.listdir', 'tf.io.gfile.listdir', ({(40, 40, 40, 49): 'input_dir'}, {}), '(input_dir)', True, 'import tensorflow as tf\n'), ((81, 4, 81, 17), 'absl.app.run', 'app.run', ({(81, 12, 81, 16): 'main'}, {}), '(main)', False, 'from absl import app\n'), ((8, 40, 8, 65), 'os.path.dirname', 'os.path.dirname', ({(8, 56, 8, 64): '__file__'}, {}), '(__file__)', False, 'import os, sys\n'), ((37, 11, 37, 37), 'os.path.exists', 'os.path.exists', ({(37, 26, 37, 36): 'output_dir'}, {}), '(output_dir)', False, 'import os, sys\n'), ((38, 8, 38, 28), 'os.mkdir', 'os.mkdir', ({(38, 17, 38, 27): 'output_dir'}, {}), '(output_dir)', False, 'import os, sys\n'), ((42, 21, 42, 54), 'os.path.join', 'os.path.join', ({(42, 34, 42, 43): 'input_dir', (42, 45, 42, 53): 'filename'}, {}), '(input_dir, filename)', False, 'import os, sys\n'), ((43, 22, 43, 56), 'os.path.join', 'os.path.join', ({(43, 35, 43, 45): 'output_dir', (43, 47, 43, 55): 'filename'}, {}), '(output_dir, filename)', False, 'import os, sys\n'), ((45, 14, 45, 36), 'PIL.Image.open', 'Image.open', ({(45, 25, 45, 35): 'input_path'}, {}), '(input_path)', False, 'from PIL import Image\n'), ((67, 19, 67, 41), 'ids.voc2012.get_colormap', 'get_voc2012_colormap', ({}, {}), '()', True, 'from ids.voc2012 import get_colormap as get_voc2012_colormap\n'), ((69, 19, 69, 44), 'ids.cityscapes_fine.get_colormap', 'get_cityscapes_colormap', ({}, {}), '()', True, 'from ids.cityscapes_fine import get_colormap as get_cityscapes_colormap\n')]