# Copyright 2021 Huawei Technologies Co., Ltd

#

# 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 sys

import os

import scipy

import shutil





def move_valimg(val_dir='./val', devkit_dir='./ILSVRC2012_devkit_t12'):

    """

    move valimg to correspongding folders.

    val_id(start from 1) -> ILSVRC_ID(start from 1) -> WIND

    organize like:

    /val

       /n01440764

           images

       /n01443537

           images

        .....

    """

    # load synset, val ground truth and val images list

    synset = scipy.io.loadmat(os.path.join(devkit_dir, 'data', 'meta.mat'))

    

    ground_truth = open(os.path.join(devkit_dir, 'data', 'ILSVRC2012_validation_ground_truth.txt'))

    lines = ground_truth.readlines()

    labels = [int(line[:-1]) for line in lines]

    

    root, _, filenames = next(os.walk(val_dir))

    for filename in filenames:

        # val image name -> ILSVRC ID -> WIND

        val_id = int(filename.split('.')[0].split('_')[-1])

        ILSVRC_ID = labels[val_id-1]

        WIND = synset['synsets'][ILSVRC_ID-1][0][1][0]

        print("val_id:%d, ILSVRC_ID:%d, WIND:%s" % (val_id, ILSVRC_ID, WIND))



        # move val images

        output_dir = os.path.join(root, WIND)

        if os.path.isdir(output_dir):

            pass

        else:

            os.mkdir(output_dir)

        shutil.move(os.path.join(root, filename), os.path.join(output_dir, filename))



def main(val_path, devkit_path):

	move_valimg(val_path, devkit_path)



if __name__ == '__main__':

	val_path = sys.argv[1]

    devkit_path = sys.argv[2]

	main(val_path, devkit_path)