# Copyright (c) 2017 Sony Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import os
from scipy.misc import imread
from args import get_args
import matplotlib.pyplot as plt


def get_color():
    # RGB format
    return np.array([[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128], [120, 0, 128], [0, 128, 128], [128, 128, 128], [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0], [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128], [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0], [0, 64, 128], [224, 224, 192]])


def encode_label(label):
    '''
    Converting pixel values to corresponding class numbers. Assuming that the input label in 3-dim(h,w,c) and in BGR fromat read from cv2
    '''

    h, w, c = label.shape
    new_label = np.zeros((h, w, 1), dtype=np.int32)

    cls_to_clr_map = get_color()

    for i in range(cls_to_clr_map.shape[0]):
        #new_label[(label == cls_to_clr_map[i])[:,:,0]] = i
        #new_label[np.argwhere((label.astype(np.int32) == cls_to_clr_map[i]).all(axis=2))]=i
        print(np.where((label.astype(np.int32) == [120, 0, 128]).all(axis=2)))
        if i == 21:
            new_label[np.where(
                (label.astype(np.int32) == cls_to_clr_map[i]).all(axis=2))] = 255
        else:
            new_label[np.where(
                (label.astype(np.int32) == cls_to_clr_map[i]).all(axis=2))] = i

    return new_label


# this method should generate train-image.txt and train-label.txt
def generate_path_files(data_dir, train_file, val_file):

    ti = open('train_image.txt', 'w')
    tl = open('train_label.txt', 'w')
    vi = open('val_image.txt', 'w')
    vl = open('val_label.txt', 'w')

    rootdir = data_dir

    train_text_file = open(train_file, "r")
    lines = [line[:-1] for line in train_text_file]
    for line in lines:
        if os.path.exists(data_dir+'JPEGImages/'+line+'.jpg'):
            ti.write(data_dir+'JPEGImages/'+line+'.jpg' + '\n')
            assert (os.path.isfile(data_dir+'SegmentationClass/encoded/'+line +
                                   '.npy')), "No matching label file for image : " + line + '.jpg'
            tl.write(data_dir+'SegmentationClass/encoded/'+line + '.npy' + '\n')

    val_text_file = open(val_file, "r")
    lines = [line[:-1] for line in val_text_file]
    for line in lines:
        if os.path.exists(data_dir+'JPEGImages/'+line+'.jpg'):
            vi.write(data_dir+'JPEGImages/'+line+'.jpg' + '\n')
            assert (os.path.isfile(data_dir+'SegmentationClass/encoded/'+line +
                                   '.npy')), "No matching label file for image : " + line + '.jpg'
            vl.write(data_dir+'SegmentationClass/encoded/'+line + '.npy' + '\n')

    ti.close()
    tl.close()
    vi.close()
    vl.close()


def main():
    '''
    Arguments:
    train-file = txt file containing randomly selected image filenames to be taken as training set.
    val-file = txt file containing randomly selected image filenames to be taken as validation set.
    data-dir = dataset directory
    Usage: python dataset_utils.py --train-file="" --val-file="" --data_dir=""
    '''

    args = get_args()
    data_dir = args.data_dir

    if not os.path.exists(data_dir+'SegmentationClass/' + 'encoded/'):
        os.makedirs(data_dir+'SegmentationClass/' + 'encoded/')
    for filename in os.listdir(data_dir+'SegmentationClass/'):
        if os.path.isdir(data_dir+'SegmentationClass/' + filename):
            continue
        label = imread(data_dir+'SegmentationClass/' +
                       filename).astype('float32')
        label = encode_label(label)
        np.save(data_dir+'SegmentationClass/' + 'encoded/' +
                filename.split('.')[0] + '.npy', label)

    generate_path_files(args.data_dir, args.train_file, args.val_file)


if __name__ == '__main__':
    main()

[((24, 11, 24, 334), 'numpy.array', 'np.array', ({(24, 20, 24, 333): '[[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128], [120, 0, \n    128], [0, 128, 128], [128, 128, 128], [64, 0, 0], [192, 0, 0], [64, 128,\n    0], [192, 128, 0], [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, \n    128, 128], [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0], [0, 64,\n    128], [224, 224, 192]]'}, {}), '([[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128],\n    [120, 0, 128], [0, 128, 128], [128, 128, 128], [64, 0, 0], [192, 0, 0],\n    [64, 128, 0], [192, 128, 0], [64, 0, 128], [192, 0, 128], [64, 128, 128\n    ], [192, 128, 128], [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0\n    ], [0, 64, 128], [224, 224, 192]])', True, 'import numpy as np\n'), ((33, 16, 33, 51), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((94, 11, 94, 21), 'args.get_args', 'get_args', ({}, {}), '()', False, 'from args import get_args\n'), ((99, 20, 99, 61), 'os.listdir', 'os.listdir', ({(99, 31, 99, 60): "(data_dir + 'SegmentationClass/')"}, {}), "(data_dir + 'SegmentationClass/')", False, 'import os\n'), ((64, 11, 64, 61), 'os.path.exists', 'os.path.exists', ({(64, 26, 64, 60): "(data_dir + 'JPEGImages/' + line + '.jpg')"}, {}), "(data_dir + 'JPEGImages/' + line + '.jpg')", False, 'import os\n'), ((73, 11, 73, 61), 'os.path.exists', 'os.path.exists', ({(73, 26, 73, 60): "(data_dir + 'JPEGImages/' + line + '.jpg')"}, {}), "(data_dir + 'JPEGImages/' + line + '.jpg')", False, 'import os\n'), ((97, 11, 97, 69), 'os.path.exists', 'os.path.exists', ({(97, 26, 97, 68): "(data_dir + 'SegmentationClass/' + 'encoded/')"}, {}), "(data_dir + 'SegmentationClass/' + 'encoded/')", False, 'import os\n'), ((98, 8, 98, 63), 'os.makedirs', 'os.makedirs', ({(98, 20, 98, 62): "(data_dir + 'SegmentationClass/' + 'encoded/')"}, {}), "(data_dir + 'SegmentationClass/' + 'encoded/')", False, 'import os\n'), ((100, 11, 100, 66), 'os.path.isdir', 'os.path.isdir', ({(100, 25, 100, 65): "(data_dir + 'SegmentationClass/' + filename)"}, {}), "(data_dir + 'SegmentationClass/' + filename)", False, 'import os\n'), ((66, 20, 67, 42), 'os.path.isfile', 'os.path.isfile', ({(66, 35, 67, 41): "(data_dir + 'SegmentationClass/encoded/' + line + '.npy')"}, {}), "(data_dir + 'SegmentationClass/encoded/' + line + '.npy')", False, 'import os\n'), ((75, 20, 76, 42), 'os.path.isfile', 'os.path.isfile', ({(75, 35, 76, 41): "(data_dir + 'SegmentationClass/encoded/' + line + '.npy')"}, {}), "(data_dir + 'SegmentationClass/encoded/' + line + '.npy')", False, 'import os\n'), ((102, 16, 103, 32), 'scipy.misc.imread', 'imread', ({(102, 23, 103, 31): "data_dir + 'SegmentationClass/' + filename"}, {}), "(data_dir + 'SegmentationClass/' + filename)", False, 'from scipy.misc import imread\n')]