# ================================================================
# 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, 'sys.path.insert', 'sys.path.insert', (['(1)', 'rootpath'], {}), False, 'import os, sys\n'), (25, 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""input_dir"""', 'None', '"""input dir path"""'], {}), False, 'from absl import flags\n'), (26, 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""output_dir"""', 'None', '"""output dir path"""'], {}), False, 'from absl import flags\n'), (27, 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""colormap"""', '"""voc2012"""', '"""colormap name"""'], {}), False, 'from absl import flags\n'), (28, 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', (['"""ignore_label"""', '(255)', '"""ignore label"""'], {}), False, 'from absl import flags\n'), (40, 'tensorflow.io.gfile.listdir', 'tf.io.gfile.listdir', (['input_dir'], {}), True, 'import tensorflow as tf\n'), (81, 'absl.app.run', 'app.run', (['main'], {}), False, 'from absl import app\n'), (8, 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), False, 'import os, sys\n'), (37, 'os.path.exists', 'os.path.exists', (['output_dir'], {}), False, 'import os, sys\n'), (38, 'os.mkdir', 'os.mkdir', (['output_dir'], {}), False, 'import os, sys\n'), (42, 'os.path.join', 'os.path.join', (['input_dir', 'filename'], {}), False, 'import os, sys\n'), (43, 'os.path.join', 'os.path.join', (['output_dir', 'filename'], {}), False, 'import os, sys\n'), (45, 'PIL.Image.open', 'Image.open', (['input_path'], {}), False, 'from PIL import Image\n'), (67, 'ids.voc2012.get_colormap', 'get_voc2012_colormap', ([], {}), True, 'from ids.voc2012 import get_colormap as get_voc2012_colormap\n'), (69, 'ids.cityscapes_fine.get_colormap', 'get_cityscapes_colormap', ([], {}), True, 'from ids.cityscapes_fine import get_colormap as get_cityscapes_colormap\n')]