Spaces:
Runtime error
Runtime error
File size: 2,000 Bytes
31ac1b2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
# DanBooru IMage Utility functions
# Taken from https://huggingface.co./spaces/SmilingWolf/wd-v1-4-tags
import cv2
import numpy as np
from PIL import Image
import PIL
def smart_imread(img, flag=cv2.IMREAD_UNCHANGED):
if img.endswith(".gif"):
img = Image.open(img)
img = img.convert("RGB")
img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
else:
img = cv2.imread(img, flag)
return img
def smart_24bit(img):
if img.dtype is np.dtype(np.uint16):
img = (img / 257).astype(np.uint8)
if len(img.shape) == 2:
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
elif img.shape[2] == 4:
trans_mask = img[:, :, 3] == 0
img[trans_mask] = [255, 255, 255, 255]
img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
return img
def make_square(img, target_size):
old_size = img.shape[:2]
desired_size = max(old_size)
desired_size = max(desired_size, target_size)
delta_w = desired_size - old_size[1]
delta_h = desired_size - old_size[0]
top, bottom = delta_h // 2, delta_h - (delta_h // 2)
left, right = delta_w // 2, delta_w - (delta_w // 2)
color = [255, 255, 255]
new_im = cv2.copyMakeBorder(
img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color
)
return new_im
def smart_resize(img, size):
# Assumes the image has already gone through make_square
if img.shape[0] > size:
img = cv2.resize(img, (size, size), interpolation=cv2.INTER_AREA)
elif img.shape[0] < size:
img = cv2.resize(img, (size, size), interpolation=cv2.INTER_CUBIC)
return img
def preprocess_image(img):
image = img.convert('RGBA')
new_image = PIL.Image.new('RGBA', image.size, 'WHITE')
new_image.paste(image, mask=image)
image = new_image.convert('RGB')
image = np.asarray(image)
image = make_square(image, 384)
image = smart_resize(image, 384)
image = image.astype(np.float32)
return Image.fromarray(np.uint8(image)) |