import os
import shutil
import numpy as np
import cv2
from PIL import Image

def cv2_imread(path):
    img = cv2.imdecode(np.fromfile(path, dtype=np.uint8), 1)
    return img

def cv2_imwrite(image, image_path, type='jpg'):
    cv2.imencode('.{}'.format(type), image)[1].tofile(image_path)

def my_letter_box(img, size=(640, 640)):  #
    h, w, c = img.shape
    r = min(size[0] / h, size[1] / w)
    new_h, new_w = int(h * r), int(w * r)
    top = int((size[0] - new_h) / 2)
    left = int((size[1] - new_w) / 2)

    bottom = size[0] - new_h - top
    right = size[1] - new_w - left
    img_resize = cv2.resize(img, (new_w, new_h))
    img = cv2.copyMakeBorder(img_resize, top, bottom, left, right, borderType=cv2.BORDER_CONSTANT,
                             value=(128, 128, 128))
    return img, r, left, top
    

        
def jpg2jpg_resize(input_path,out_dir):
    try:
        img = cv2_imread(input_path)
        nw = 640
        nh = 640
        #img = img[:, :, ::-1] #BGR2RGB
        img = cv2.resize(img, (nw, nh))
        #result = np.transpose(img, [2,0,1])
        #output_name = input_path.rsplit('.', 1)[0] + ".bgr"
        result = img
        output_name = out_dir +input_path.split("/")[-1].rsplit('.', 1)[0] + ".jpg"
        result.tofile(output_name)
    except Exception as except_err:
        print(except_err)
        return 1

    else:
        return 0

def jpg2bgr_pad_letterbox(input_path,out_dir):
    try:
        img = cv2_imread(input_path)
        nw = 640
        nh = 640
        
        img_0, r, left, top = my_letter_box(img, size=(nw, nh))
        output_name = out_dir +input_path.split("/")[-1].rsplit('.', 1)[0] + ".jpg"
        result = img_0
        #result.tofile(output_name)
        cv2_imwrite(img_0,output_name)
    except Exception as except_err:
        print(except_err)
        return 1

    else:
        return 0
        
def jpg2bgr_pad_no_letterbox(input_path,out_dir):
    try:
        img = cv2_imread(input_path)
        nw = 640
        nh = 640
        #img = img[:, :, ::-1] #BGR2RGB
        img = cv2.resize(img, (nw, nh))
        result = np.transpose(img, [2,0,1])
        #output_name = input_path.rsplit('.', 1)[0] + ".bgr"
        output_name = out_dir +input_path.split("/")[-1].rsplit('.', 1)[0] + ".bgr"
        result.tofile(output_name)
    except Exception as except_err:
        print(except_err)
        return 1

    else:
        return 0

def jpg2bgr_pad(input_path,out_dir):
    try:
        img = cv2_imread(input_path)

        img_0, r, left, top = my_letter_box(img, size=(640, 640))
        result = np.transpose(img_0, [2,0,1])
        #output_name = input_path.rsplit('.', 1)[0] + ".bgr"
        output_name = out_dir +input_path.split("/")[-1].rsplit('.', 1)[0] + ".bgr"
        result.tofile(output_name)
    except Exception as except_err:
        print(except_err)
        return 1

    else:
        return 0

def jpg2bgr_pad_PIL(input_path,out_dir):
    try:
        #img = cv2_imread(input_path)
        img = Image.open(input_path)
        nw = 480
        nh = 480
        img_0 = img.resize((nw, nh), Image.BICUBIC)
        img_0 = np.array(img_0)
        result = np.transpose(img_0, [2,0,1])
        #output_name = input_path.rsplit('.', 1)[0] + ".bgr"
        output_name = out_dir +input_path.split("/")[-1].rsplit('.', 1)[0] + ".bgr"
        print(output_name)
        result.tofile(output_name)
    except Exception as except_err:
        print(except_err)
        return 1

    else:
        return 0
def bgr2nv21(input_path,out_dir):
    try:
        bgr = cv2.imread(input_path)
        bgr = cv2.resize(bgr, (480, 480))
        i420 = cv2.cvtColor(bgr, cv2.COLOR_BGR2YUV_I420)
        height = bgr.shape[0]
        width = bgr.shape[1]
    
        u = i420[height: height + height // 4, :]
        u = u.reshape((1, height // 4 * width))
        v = i420[height + height // 4: height + height // 2, :]
        v = v.reshape((1, height // 4 * width))
        uv = np.zeros((1, height // 4 * width * 2))
        uv[:, 0::2] = v
        uv[:, 1::2] = u
        uv = uv.reshape((height // 2, width))
        nv21 = np.zeros((height + height // 2, width))
        nv21[0:height, :] = i420[0:height, :]
        nv21[height::, :] = uv
        #output_name = input_path.rsplit('.', 1)[0] + ".yuv"
        output_name = out_dir +input_path.split("/")[-1].rsplit('.', 1)[0] + ".yuv"
        nv21.astype("int8").tofile(output_name)
    except Exception as except_err:
        print(except_err)
        return 1
    else:
        return 0

if __name__ == "__main__":
    count_ok = 0
    count_ng = 0
    images = os.listdir(r'/data//20240927/')
    dir = os.path.realpath(r"/data/20240927/")
    out_dir = r"/data/HI3516/t20240927_out_jpg/"
    for image_name in images:
        if not (image_name.lower().endswith((".bmp", ".dib", ".jpeg", ".jpg", ".jpe",
        ".png", ".pbm", ".pgm", ".ppm", ".sr", ".ras", ".tiff", ".tif"))):
            continue
        print("start to process image {}....".format(image_name))
        image_path = os.path.join(dir, image_name)
        ret = jpg2bgr_pad_letterbox(image_path,out_dir)
        if ret == 0:
            print("process image {} successfully".format(image_name))
            count_ok = count_ok + 1
        elif ret == 1:
            print("failed to process image {}".format(image_name))
            count_ng = count_ng + 1
    print("{} images in total, {} images process successfully, {} images process failed"
          .format(count_ok + count_ng, count_ok, count_ng))