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

import numpy as np

import torch



def preprocess(file_path, bin_path):

    in_files = os.listdir(file_path)

    if not os.path.exists(bin_path):

        os.makedirs(bin_path)

    i = 0

    for file in sorted(in_files):   

        i = i + 1

        print(file, "===", i)

        image = cv2.imread(os.path.join(file_path, file))

        height, width = image.shape[0:2]

        new_height = int(height)

        new_width  = int(width)

        #Fix size testing

        inp_height, inp_width = 512, 512

        c = np.array([new_width / 2., new_height / 2.], dtype=np.float32)

        s = max(height, width) * 1.0

        trans_input = get_affine_transform(c, s, 0, [inp_width, inp_height])

        resized_image = cv2.resize(image, (new_width, new_height))

        inp_image = cv2.warpAffine(

        resized_image, trans_input, (inp_width, inp_height),flags=cv2.INTER_LINEAR)

        

        inp_image = ((inp_image / 255. - [[[0.40789655,0.44719303,0.47026116]]]) / [[[0.2886383,0.27408165,0.27809834]]]).astype(np.float32)

        images = inp_image.transpose(2, 0, 1).reshape(1, 3, inp_height, inp_width)

        images = torch.from_numpy(images) 

        img = np.array(images).astype(np.float32)

        img.tofile(os.path.join(bin_path, file.split('.')[0] + '.bin'))



def get_affine_transform(center,

                         scale,

                         rot,

                         output_size,

                         shift=np.array([0, 0], dtype=np.float32),

                         inv=0):

    if not isinstance(scale, np.ndarray) and not isinstance(scale, list):

        scale = np.array([scale, scale], dtype=np.float32)



    scale_tmp = scale

    src_w = scale_tmp[0]

    dst_w = output_size[0]

    dst_h = output_size[1]



    rot_rad = np.pi * rot / 180

    src_dir = get_dir([0, src_w * -0.5], rot_rad)

    dst_dir = np.array([0, dst_w * -0.5], np.float32)



    src = np.zeros((3, 2), dtype=np.float32)

    dst = np.zeros((3, 2), dtype=np.float32)

    src[0, :] = center + scale_tmp * shift

    src[1, :] = center + src_dir + scale_tmp * shift

    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]

    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5], np.float32) + dst_dir



    src[2:, :] = get_3rd_point(src[0, :], src[1, :])

    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])



    if inv:

        trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))

    else:

        trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))



    return trans



def get_3rd_point(a, b):

    direct = a - b

    return b + np.array([-direct[1], direct[0]], dtype=np.float32)



def get_dir(src_point, rot_rad):

    sn, cs = np.sin(rot_rad), np.cos(rot_rad)



    src_result = [0, 0]

    src_result[0] = src_point[0] * cs - src_point[1] * sn

    src_result[1] = src_point[0] * sn + src_point[1] * cs



    return src_result

if __name__ == "__main__":

    file_path = os.path.abspath(sys.argv[1])

    bin_path = os.path.abspath(sys.argv[2])

    preprocess(file_path, bin_path)