CheckCardBoxAndDefectServer/app/utils/card_inference/predict_preprocess.py

import copy
import random
import numpy as np
import math
import cv2
import numpy as np
import math

import torch
import torch.nn as nn
import torchvision.transforms as transforms
import cv2
import __main__
import sys
import os
import time


from app.utils.card_inference.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