# ================================================================
# 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)

[((326, 354), 'sys.path.insert', 'sys.path.insert', (['(1)', 'rootpath'], {}), '(1, rootpath)\n', (341, 354), False, 'import os, sys\n'), ((634, 690), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""input_dir"""', 'None', '"""input dir path"""'], {}), "('input_dir', None, 'input dir path')\n", (653, 690), False, 'from absl import flags\n'), ((691, 749), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""output_dir"""', 'None', '"""output dir path"""'], {}), "('output_dir', None, 'output dir path')\n", (710, 749), False, 'from absl import flags\n'), ((750, 809), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""colormap"""', '"""voc2012"""', '"""colormap name"""'], {}), "('colormap', 'voc2012', 'colormap name')\n", (769, 809), False, 'from absl import flags\n'), ((810, 867), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', (['"""ignore_label"""', '(255)', '"""ignore label"""'], {}), "('ignore_label', 255, 'ignore label')\n", (830, 867), False, 'from absl import flags\n'), ((1089, 1119), 'tensorflow.io.gfile.listdir', 'tf.io.gfile.listdir', (['input_dir'], {}), '(input_dir)\n', (1108, 1119), True, 'import tensorflow as tf\n'), ((2042, 2055), 'absl.app.run', 'app.run', (['main'], {}), '(main)\n', (2049, 2055), False, 'from absl import app\n'), ((275, 300), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (290, 300), False, 'import os, sys\n'), ((1011, 1037), 'os.path.exists', 'os.path.exists', (['output_dir'], {}), '(output_dir)\n', (1025, 1037), False, 'import os, sys\n'), ((1047, 1067), 'os.mkdir', 'os.mkdir', (['output_dir'], {}), '(output_dir)\n', (1055, 1067), False, 'import os, sys\n'), ((1143, 1176), 'os.path.join', 'os.path.join', (['input_dir', 'filename'], {}), '(input_dir, filename)\n', (1155, 1176), False, 'import os, sys\n'), ((1199, 1233), 'os.path.join', 'os.path.join', (['output_dir', 'filename'], {}), '(output_dir, filename)\n', (1211, 1233), False, 'import os, sys\n'), ((1249, 1271), 'PIL.Image.open', 'Image.open', (['input_path'], {}), '(input_path)\n', (1259, 1271), False, 'from PIL import Image\n'), ((1676, 1698), 'ids.voc2012.get_colormap', 'get_voc2012_colormap', ([], {}), '()\n', (1696, 1698), True, 'from ids.voc2012 import get_colormap as get_voc2012_colormap\n'), ((1758, 1783), 'ids.cityscapes_fine.get_colormap', 'get_cityscapes_colormap', ([], {}), '()\n', (1781, 1783), True, 'from ids.cityscapes_fine import get_colormap as get_cityscapes_colormap\n')]