Views
No views yet
1import torch
2from transformers import CLIPProcessor, CLIPModel
3
4IMG_SIZE = 224
5DEVICE = "cpu" # can also be "cuda" or "mps"
6LOGIT_SCALE = 100 # based on OpenAI's CLIP example code
7NORMALIZE_SCORING = True
8
9model_path="uiclip_jitteredwebsites-2-224-paraphrased_webpairs" # can also be regular or human pairs variants
10processor_path="openai/clip-vit-base-patch32"
11
12model = CLIPModel.from_pretrained(model_path)
13model = model.eval()
14model = model.to(DEVICE)
15
16processor = CLIPProcessor.from_pretrained(processor_path)
17
18def compute_quality_scores(input_list):
19 # input_list is a list of types where the first element is a description and the second is a PIL image
20 description_list = ["ui screenshot. well-designed. " + input_item[0] for input_item in input_list]
21 img_list = [input_item[1] for input_item in input_list]
22 text_embeddings_tensor = compute_description_embeddings(description_list) # B x H
23 img_embeddings_tensor = compute_image_embeddings(img_list) # B x H
24
25 # normalize tensors
26 text_embeddings_tensor /= text_embeddings_tensor.norm(dim=-1, keepdim=True)
27 img_embeddings_tensor /= img_embeddings_tensor.norm(dim=-1, keepdim=True)
28
29 if NORMALIZE_SCORING:
30 text_embeddings_tensor_poor = compute_description_embeddings([d.replace("well-designed. ", "poor design. ") for d in description_list]) # B x H
31 text_embeddings_tensor_poor /= text_embeddings_tensor_poor.norm(dim=-1, keepdim=True)
32 text_embeddings_tensor_all = torch.stack((text_embeddings_tensor, text_embeddings_tensor_poor), dim=1) # B x 2 x H
33 else:
34 text_embeddings_tensor_all = text_embeddings_tensor.unsqueeze(1)
35
36 img_embeddings_tensor = img_embeddings_tensor.unsqueeze(1) # B x 1 x H
37
38 scores = (LOGIT_SCALE * img_embeddings_tensor @ text_embeddings_tensor_all.permute(0, 2, 1)).squeeze(1)
39
40 if NORMALIZE_SCORING:
41 scores = scores.softmax(dim=-1)
42
43 return scores[:, 0]
44
45def compute_description_embeddings(descriptions):
46 inputs = processor(text=descriptions, return_tensors="pt", padding=True)
47 inputs['input_ids'] = inputs['input_ids'].to(DEVICE)
48 inputs['attention_mask'] = inputs['attention_mask'].to(DEVICE)
49 text_embedding = model.get_text_features(**inputs)
50 return text_embedding
51
52def compute_image_embeddings(image_list):
53 windowed_batch = [slide_window_over_image(img, IMG_SIZE) for img in image_list]
54 inds = []
55 for imgi in range(len(windowed_batch)):
56 inds.append([imgi for _ in windowed_batch[imgi]])
57
58 processed_batch = [item for sublist in windowed_batch for item in sublist]
59 inputs = processor(images=processed_batch, return_tensors="pt")
60 # run all sub windows of all images in batch through the model
61 inputs['pixel_values'] = inputs['pixel_values'].to(DEVICE)
62 with torch.no_grad():
63 image_features = model.get_image_features(**inputs)
64
65 # output contains all subwindows, need to mask for each image
66 processed_batch_inds = torch.tensor([item for sublist in inds for item in sublist]).long().to(image_features.device)
67 embed_list = []
68 for i in range(len(image_list)):
69 mask = processed_batch_inds == i
70 embed_list.append(image_features[mask].mean(dim=0))
71 image_embedding = torch.stack(embed_list, dim=0)
72 return image_embedding
73
74def preresize_image(image, image_size):
75 aspect_ratio = image.width / image.height
76 if aspect_ratio > 1:
77 image = image.resize((int(aspect_ratio * image_size), image_size))
78 else:
79 image = image.resize((image_size, int(image_size / aspect_ratio)))
80 return image
81
82def slide_window_over_image(input_image, img_size):
83 input_image = preresize_image(input_image, img_size)
84 width, height = input_image.size
85 square_size = min(width, height)
86 longer_dimension = max(width, height)
87 num_steps = (longer_dimension + square_size - 1) // square_size
88
89 if num_steps > 1:
90 step_size = (longer_dimension - square_size) // (num_steps - 1)
91 else:
92 step_size = square_size
93
94 cropped_images = []
95
96 for y in range(0, height - square_size + 1, step_size if height > width else square_size):
97 for x in range(0, width - square_size + 1, step_size if width > height else square_size):
98 left = x
99 upper = y
100 right = x + square_size
101 lower = y + square_size
102 cropped_image = input_image.crop((left, upper, right, lower))
103 cropped_images.append(cropped_image)
104
105 return cropped_images
106
107
108# compute the quality scores for a list of descriptions (strings) and images (PIL images)
109prediction_scores = compute_quality_scores(list(zip(test_descriptions, test_images)))