sd/prompt_helper/helper.py

import onnx
import onnxruntime as ort
import numpy as np

VIT_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-tagger-v3"


def preprocess_image(img):
    bgr_img = np.array(img)[:, :, ::-1].copy()
    size = max(bgr_img.shape[0:2])
    pad_x = size - bgr_img.shape[1]
    pad_y = size - bgr_img.shape[0]
    pad_l = pad_x // 2
    pad_t = pad_y // 2

    #add paddings to squaring image
    np.pad(bgr_img, ((pad_t, pad_y - pad_t), (pad_l, pad_x - pad_l), (0, 0)), mode="constant", constant_values=255)

    #adaptive resize
    interp = cv2.INTER_AREA if size > IMAGE_SIZE else cv2.INTER_LANCZOS4
    bgr_img = cv2.resize(bgr_img, (IMAGE_SIZE, IMAGE_SIZE), interpolation=interp)

    bgr_img = bgr_img.astype(np.float32)