import math
import cv2
import numpy as np
import math
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import cv2

from app.utils.data_augmentation import LetterBox


def predict_preprocess(img_bgr, imgSize_train):
    device_str = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

    letterBox = LetterBox(imgSize_train)
    img_rgb_letterbox = letterBox.handle_single_img(img_rgb)

    img_np = np.array(img_rgb_letterbox)

    imgTensor = torch.tensor(img_np, dtype=torch.float32, device=device_str)

    # 将所有元素值除以255，进行归一化
    imgTensor = imgTensor * (1 / 255.0)

    # 把形状从[H, W, C] 改为 [C, H, W]
    imgTensor_CHW = imgTensor.permute(2, 0, 1)

    normaliz_operate_c3 = transforms.Compose([
        transforms.Normalize(mean=(0, 0, 0), std=(1, 1, 1)),
    ])
    imgTensor_CHW_norm = normaliz_operate_c3(imgTensor_CHW)

    return imgTensor_CHW_norm