File size: 3,910 Bytes
f961e67 |
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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import torch
import numpy as np
from torchvision import transforms
from PIL import Image, ImageFilter
import random
def set_random_seed(seed=0):
torch.manual_seed(seed + 0)
torch.cuda.manual_seed(seed + 1)
torch.cuda.manual_seed_all(seed + 2)
np.random.seed(seed + 3)
torch.cuda.manual_seed_all(seed + 4)
random.seed(seed + 5)
def transform_img(image, target_size=512):
tform = transforms.Compose(
[
transforms.Resize(target_size),
transforms.CenterCrop(target_size),
transforms.ToTensor(),
]
)
image = tform(image)
return 2.0 * image - 1.0
def latents_to_imgs(pipe, latents):
x = pipe.decode_image(latents)
x = pipe.torch_to_numpy(x)
x = pipe.numpy_to_pil(x)
return x
def image_distortion(img,
seed: int = 42,
random_crop_ratio: float = None,
random_drop_ratio: float = None,
resize_ratio: float = None,
gaussian_blur_r: int = None, #
gaussian_std: float = None,
sp_prob: float = None):
if random_crop_ratio is not None:
set_random_seed(seed)
width, height, c = np.array(img).shape
img = np.array(img)
new_width = int(width * random_crop_ratio)
new_height = int(height * random_crop_ratio)
start_x = np.random.randint(0, width - new_width + 1)
start_y = np.random.randint(0, height - new_height + 1)
end_x = start_x + new_width
end_y = start_y + new_height
padded_image = np.zeros_like(img)
padded_image[start_y:end_y, start_x:end_x] = img[start_y:end_y, start_x:end_x]
img = Image.fromarray(padded_image)
if random_drop_ratio is not None:
set_random_seed(seed)
width, height, c = np.array(img).shape
img = np.array(img)
new_width = int(width * random_drop_ratio)
new_height = int(height * random_drop_ratio)
start_x = np.random.randint(0, width - new_width + 1)
start_y = np.random.randint(0, height - new_height + 1)
padded_image = np.zeros_like(img[start_y:start_y + new_height, start_x:start_x + new_width])
img[start_y:start_y + new_height, start_x:start_x + new_width] = padded_image
img = Image.fromarray(img)
if resize_ratio is not None:
img_shape = np.array(img).shape
resize_size = int(img_shape[0] * resize_ratio)
img = transforms.Resize(size=resize_size)(img)
img = transforms.Resize(size=img_shape[0])(img)
if gaussian_blur_r is not None:
img = img.filter(ImageFilter.GaussianBlur(radius=gaussian_blur_r))
if gaussian_std is not None:
img_shape = np.array(img).shape
g_noise = np.random.normal(0, gaussian_std, img_shape) * 255
g_noise = g_noise.astype(np.uint8)
img = Image.fromarray(np.clip(np.array(img) + g_noise, 0, 255))
if sp_prob is not None:
c,h,w = np.array(img).shape
prob_zero = sp_prob / 2
prob_one = 1 - prob_zero
rdn = np.random.rand(c,h,w)
img = np.where(rdn > prob_one, np.zeros_like(img), img)
img = np.where(rdn < prob_zero, np.ones_like(img)*255, img)
img = Image.fromarray(img)
return img
def measure_similarity(images, prompt, model, clip_preprocess, tokenizer, device):
with torch.no_grad():
img_batch = [clip_preprocess(i).unsqueeze(0) for i in images]
img_batch = torch.concatenate(img_batch).to(device)
image_features = model.encode_image(img_batch)
text = tokenizer([prompt]).to(device)
text_features = model.encode_text(text)
image_features /= image_features.norm(dim=-1, keepdim=True)
text_features /= text_features.norm(dim=-1, keepdim=True)
return (image_features @ text_features.T).mean(-1)
|