image-evaluator 0.0.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- image_evaluator-0.0.1/LICENSE +21 -0
- image_evaluator-0.0.1/PKG-INFO +28 -0
- image_evaluator-0.0.1/README.md +2 -0
- image_evaluator-0.0.1/image_evaluator/__init__.py +0 -0
- image_evaluator-0.0.1/image_evaluator/arcface_dist_predictor.py +90 -0
- image_evaluator-0.0.1/image_evaluator/clip_score_predictor.py +296 -0
- image_evaluator-0.0.1/image_evaluator/laion_ai_aesthetic_predictor.py +100 -0
- image_evaluator-0.0.1/image_evaluator/main.py +67 -0
- image_evaluator-0.0.1/pyproject.toml +31 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 Faych Chen
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
Metadata-Version: 2.3
|
|
2
|
+
Name: image-evaluator
|
|
3
|
+
Version: 0.0.1
|
|
4
|
+
Summary: An automatic image evaluation script
|
|
5
|
+
License: MIT
|
|
6
|
+
Author: neverbiasu
|
|
7
|
+
Requires-Python: >=3.8,<4.0
|
|
8
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
9
|
+
Classifier: Programming Language :: Python :: 3
|
|
10
|
+
Classifier: Programming Language :: Python :: 3.8
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
16
|
+
Requires-Dist: clip-score
|
|
17
|
+
Requires-Dist: insightface
|
|
18
|
+
Requires-Dist: numpy
|
|
19
|
+
Requires-Dist: open-clip-torch
|
|
20
|
+
Requires-Dist: pillow
|
|
21
|
+
Requires-Dist: torch (>=1.7.0)
|
|
22
|
+
Requires-Dist: tqdm
|
|
23
|
+
Requires-Dist: transformers
|
|
24
|
+
Description-Content-Type: text/markdown
|
|
25
|
+
|
|
26
|
+
# image-evaluator
|
|
27
|
+
An automatic image evaluation script
|
|
28
|
+
|
|
File without changes
|
|
@@ -0,0 +1,90 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import torch
|
|
3
|
+
import torch.nn.functional as F
|
|
4
|
+
from torchvision import transforms
|
|
5
|
+
from PIL import Image
|
|
6
|
+
from insightface.app import FaceAnalysis
|
|
7
|
+
import numpy as np
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class ArcFaceDistPredictor:
|
|
11
|
+
def __init__(self, model_name="buffalo_l", device=None):
|
|
12
|
+
"""Initialize ArcFace distance predictor"""
|
|
13
|
+
if device is None:
|
|
14
|
+
ctx_id = 0 if torch.cuda.is_available() else -1
|
|
15
|
+
else:
|
|
16
|
+
ctx_id = 0 if device == 'cuda' else -1
|
|
17
|
+
|
|
18
|
+
# Initialize ArcFace model
|
|
19
|
+
self.app = FaceAnalysis(model_name)
|
|
20
|
+
self.app.prepare(ctx_id=ctx_id)
|
|
21
|
+
|
|
22
|
+
# Image preprocessing
|
|
23
|
+
self.transform = transforms.Compose(
|
|
24
|
+
[
|
|
25
|
+
transforms.Resize((112, 112)),
|
|
26
|
+
transforms.ToTensor(),
|
|
27
|
+
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
|
|
28
|
+
]
|
|
29
|
+
)
|
|
30
|
+
|
|
31
|
+
def get_face_embedding(self, image_path):
|
|
32
|
+
"""Get face embedding vector
|
|
33
|
+
|
|
34
|
+
Args:
|
|
35
|
+
image_path: Image file path
|
|
36
|
+
|
|
37
|
+
Returns:
|
|
38
|
+
numpy.ndarray: Face embedding vector or None (if no face is detected)
|
|
39
|
+
"""
|
|
40
|
+
# Read image and convert to NumPy array
|
|
41
|
+
img = Image.open(image_path).convert("RGB")
|
|
42
|
+
img = np.array(img)
|
|
43
|
+
|
|
44
|
+
# Get face embedding
|
|
45
|
+
faces = self.app.get(img)
|
|
46
|
+
if len(faces) == 0:
|
|
47
|
+
return None
|
|
48
|
+
return faces[0].embedding
|
|
49
|
+
|
|
50
|
+
def evaluate_arcface_distance(self, reference_path, generated_path):
|
|
51
|
+
"""Evaluate ArcFace distance between two images
|
|
52
|
+
|
|
53
|
+
Args:
|
|
54
|
+
reference_path: Reference image file path
|
|
55
|
+
generated_path: Generated image file path
|
|
56
|
+
|
|
57
|
+
Returns:
|
|
58
|
+
float: ArcFace distance score or None (if face not detected in either image)
|
|
59
|
+
"""
|
|
60
|
+
ref_embedding = self.get_face_embedding(reference_path)
|
|
61
|
+
gen_embedding = self.get_face_embedding(generated_path)
|
|
62
|
+
if ref_embedding is None or gen_embedding is None:
|
|
63
|
+
return None
|
|
64
|
+
return (
|
|
65
|
+
1
|
|
66
|
+
- F.cosine_similarity(
|
|
67
|
+
torch.tensor(ref_embedding), torch.tensor(gen_embedding), dim=0
|
|
68
|
+
).item()
|
|
69
|
+
)
|
|
70
|
+
|
|
71
|
+
def evaluate_folder_arcface_distance(self, reference_folder, generated_folder):
|
|
72
|
+
"""Evaluate average ArcFace distance between images in two folders
|
|
73
|
+
|
|
74
|
+
Args:
|
|
75
|
+
reference_folder: Folder path containing reference images
|
|
76
|
+
generated_folder: Folder path containing generated images
|
|
77
|
+
|
|
78
|
+
Returns:
|
|
79
|
+
float: Average ArcFace distance
|
|
80
|
+
"""
|
|
81
|
+
reference_images = sorted(os.listdir(reference_folder))
|
|
82
|
+
generated_images = sorted(os.listdir(generated_folder))
|
|
83
|
+
distances = []
|
|
84
|
+
for ref_img, gen_img in zip(reference_images, generated_images):
|
|
85
|
+
ref_path = os.path.join(reference_folder, ref_img)
|
|
86
|
+
gen_path = os.path.join(generated_folder, gen_img)
|
|
87
|
+
dist = self.evaluate_arcface_distance(ref_path, gen_path)
|
|
88
|
+
if dist is not None:
|
|
89
|
+
distances.append(dist)
|
|
90
|
+
return np.mean(distances) if distances else None
|
|
@@ -0,0 +1,296 @@
|
|
|
1
|
+
# Taken from https://github.com/Taited/clip-score/blob/master/src/clip_score/clip_score.py
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
import torch
|
|
5
|
+
import os.path as osp
|
|
6
|
+
from PIL import Image
|
|
7
|
+
from tqdm import tqdm
|
|
8
|
+
from torch.utils.data import DataLoader, Dataset
|
|
9
|
+
from transformers import AutoModel, AutoProcessor, AutoTokenizer
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class DummyDataset(Dataset):
|
|
13
|
+
|
|
14
|
+
FLAGS = ['img', 'txt']
|
|
15
|
+
|
|
16
|
+
def __init__(
|
|
17
|
+
self,
|
|
18
|
+
real_path,
|
|
19
|
+
fake_path,
|
|
20
|
+
real_flag: str = 'img',
|
|
21
|
+
fake_flag: str = 'txt',
|
|
22
|
+
transform=None,
|
|
23
|
+
tokenizer=None,
|
|
24
|
+
) -> None:
|
|
25
|
+
super().__init__()
|
|
26
|
+
if real_flag not in self.FLAGS or fake_flag not in self.FLAGS:
|
|
27
|
+
raise TypeError(
|
|
28
|
+
'CLIP Score only support modality of {}. '
|
|
29
|
+
'However, get {} and {}'.format(self.FLAGS, real_flag, fake_flag)
|
|
30
|
+
)
|
|
31
|
+
self.real_folder = self._combine_without_prefix(real_path)
|
|
32
|
+
self.real_flag = real_flag
|
|
33
|
+
self.fake_folder = self._combine_without_prefix(fake_path)
|
|
34
|
+
self.fake_flag = fake_flag
|
|
35
|
+
self.transform = transform
|
|
36
|
+
self.tokenizer = tokenizer
|
|
37
|
+
# assert self._check()
|
|
38
|
+
|
|
39
|
+
def __len__(self):
|
|
40
|
+
if isinstance(self.real_folder, list):
|
|
41
|
+
real_folder_length = len(self.real_folder)
|
|
42
|
+
else:
|
|
43
|
+
real_folder_lenghth = 1
|
|
44
|
+
if isinstance(self.fake_folder, list):
|
|
45
|
+
fake_folder_length = len(self.fake_folder)
|
|
46
|
+
else:
|
|
47
|
+
fake_folder_lenghth = 1
|
|
48
|
+
return max(real_folder_lenghth, fake_folder_length)
|
|
49
|
+
|
|
50
|
+
def __getitem__(self, index):
|
|
51
|
+
if index >= len(self):
|
|
52
|
+
raise IndexError
|
|
53
|
+
if isinstance(self.real_folder, list):
|
|
54
|
+
real_path = self.real_folder[index]
|
|
55
|
+
else:
|
|
56
|
+
real_path = self.real_folder
|
|
57
|
+
if isinstance(self.fake_folder, list):
|
|
58
|
+
fake_path = self.fake_folder[index]
|
|
59
|
+
else:
|
|
60
|
+
fake_path = self.fake_folder
|
|
61
|
+
real_data = self._load_modality(real_path, self.real_flag)
|
|
62
|
+
fake_data = self._load_modality(fake_path, self.fake_flag)
|
|
63
|
+
|
|
64
|
+
sample = dict(real=real_data, fake=fake_data)
|
|
65
|
+
return sample
|
|
66
|
+
|
|
67
|
+
def _load_modality(self, path, modality):
|
|
68
|
+
if modality == 'img':
|
|
69
|
+
data = self._load_img(path)
|
|
70
|
+
elif modality == 'txt':
|
|
71
|
+
data = self._load_txt(path)
|
|
72
|
+
else:
|
|
73
|
+
raise TypeError('Got unexpected modality: {}'.format(modality))
|
|
74
|
+
return data
|
|
75
|
+
|
|
76
|
+
def _load_img(self, path):
|
|
77
|
+
img = Image.open(path)
|
|
78
|
+
if self.transform is not None:
|
|
79
|
+
img = self.transform(text=None, images=img)
|
|
80
|
+
img['pixel_values'] = img['pixel_values'][0]
|
|
81
|
+
return img
|
|
82
|
+
|
|
83
|
+
def _load_txt(self, path):
|
|
84
|
+
if osp.exists(path):
|
|
85
|
+
with open(path, 'r') as fp:
|
|
86
|
+
data = fp.read()
|
|
87
|
+
fp.close()
|
|
88
|
+
else:
|
|
89
|
+
data = path
|
|
90
|
+
if self.transform is not None:
|
|
91
|
+
data = self.tokenizer(data, padding=True, return_tensors='pt')
|
|
92
|
+
for key in data:
|
|
93
|
+
data[key] = data[key].squeeze()
|
|
94
|
+
return data
|
|
95
|
+
|
|
96
|
+
def _check(self):
|
|
97
|
+
for idx in range(len(self)):
|
|
98
|
+
real_name = self.real_folder[idx].split('.')
|
|
99
|
+
fake_name = self.fake_folder[idx].split('.')
|
|
100
|
+
if fake_name != real_name:
|
|
101
|
+
return False
|
|
102
|
+
return True
|
|
103
|
+
|
|
104
|
+
def _combine_without_prefix(self, folder_path, prefix='.'):
|
|
105
|
+
if not osp.exists(folder_path):
|
|
106
|
+
return folder_path
|
|
107
|
+
folder = []
|
|
108
|
+
for name in os.listdir(folder_path):
|
|
109
|
+
if name[0] == prefix:
|
|
110
|
+
continue
|
|
111
|
+
folder.append(osp.join(folder_path, name))
|
|
112
|
+
folder.sort()
|
|
113
|
+
return folder
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
class ClipScorePredictor:
|
|
117
|
+
def __init__(self, clip_model='openai/clip-vit-base-patch32', device=None):
|
|
118
|
+
if device is None:
|
|
119
|
+
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
120
|
+
else:
|
|
121
|
+
self.device = torch.device(device)
|
|
122
|
+
|
|
123
|
+
print('Loading CLIP model: {}'.format(clip_model))
|
|
124
|
+
self.model = AutoModel.from_pretrained(clip_model).to(self.device)
|
|
125
|
+
self.processor = AutoProcessor.from_pretrained(clip_model)
|
|
126
|
+
self.tokenizer = AutoTokenizer.from_pretrained(clip_model)
|
|
127
|
+
|
|
128
|
+
def evaluate_clip_score(
|
|
129
|
+
self,
|
|
130
|
+
real_path,
|
|
131
|
+
fake_path,
|
|
132
|
+
real_flag='img',
|
|
133
|
+
fake_flag='txt',
|
|
134
|
+
batch_size=50,
|
|
135
|
+
num_workers=None,
|
|
136
|
+
):
|
|
137
|
+
"""Evaluate CLIP score between images and text
|
|
138
|
+
|
|
139
|
+
Supports both single files and folders evaluation.
|
|
140
|
+
|
|
141
|
+
Args:
|
|
142
|
+
real_path: Path to real image or folder
|
|
143
|
+
fake_path: Path to text prompt or folder, or text string directly
|
|
144
|
+
real_flag: Type of real input modality, 'img' or 'txt'
|
|
145
|
+
fake_flag: Type of fake input modality, 'img' or 'txt'
|
|
146
|
+
batch_size: Batch size
|
|
147
|
+
num_workers: Number of workers for data loader
|
|
148
|
+
|
|
149
|
+
Returns:
|
|
150
|
+
float: CLIP score
|
|
151
|
+
"""
|
|
152
|
+
# Check if it's a single file
|
|
153
|
+
if os.path.isfile(real_path) and (
|
|
154
|
+
not os.path.exists(fake_path) or os.path.isfile(fake_path)
|
|
155
|
+
):
|
|
156
|
+
return self._evaluate_single_file(
|
|
157
|
+
real_path, fake_path, real_flag, fake_flag
|
|
158
|
+
)
|
|
159
|
+
else:
|
|
160
|
+
return self.evaluate_folder_clip_score(
|
|
161
|
+
real_path, fake_path, real_flag, fake_flag, batch_size, num_workers
|
|
162
|
+
)
|
|
163
|
+
|
|
164
|
+
def evaluate_folder_clip_score(
|
|
165
|
+
self,
|
|
166
|
+
real_path,
|
|
167
|
+
fake_path,
|
|
168
|
+
real_flag='img',
|
|
169
|
+
fake_flag='txt',
|
|
170
|
+
batch_size=50,
|
|
171
|
+
num_workers=None,
|
|
172
|
+
):
|
|
173
|
+
"""Evaluate CLIP score between multiple files in folders
|
|
174
|
+
|
|
175
|
+
Args:
|
|
176
|
+
real_path: Path to folder containing real inputs
|
|
177
|
+
fake_path: Path to folder containing fake inputs
|
|
178
|
+
real_flag: Type of real input modality, 'img' or 'txt'
|
|
179
|
+
fake_flag: Type of fake input modality, 'img' or 'txt'
|
|
180
|
+
batch_size: Batch size
|
|
181
|
+
num_workers: Number of workers for data loader
|
|
182
|
+
|
|
183
|
+
Returns:
|
|
184
|
+
float: CLIP score
|
|
185
|
+
"""
|
|
186
|
+
if num_workers is None:
|
|
187
|
+
try:
|
|
188
|
+
num_cpus = len(os.sched_getaffinity(0))
|
|
189
|
+
except AttributeError:
|
|
190
|
+
num_cpus = os.cpu_count()
|
|
191
|
+
num_workers = min(num_cpus, 8) if num_cpus is not None else 0
|
|
192
|
+
|
|
193
|
+
dataset = DummyDataset(
|
|
194
|
+
real_path,
|
|
195
|
+
fake_path,
|
|
196
|
+
real_flag,
|
|
197
|
+
fake_flag,
|
|
198
|
+
transform=self.processor,
|
|
199
|
+
tokenizer=self.tokenizer,
|
|
200
|
+
)
|
|
201
|
+
dataloader = DataLoader(
|
|
202
|
+
dataset, batch_size, num_workers=num_workers, pin_memory=True
|
|
203
|
+
)
|
|
204
|
+
|
|
205
|
+
print('Calculating CLIP Score:')
|
|
206
|
+
score_acc = 0.0
|
|
207
|
+
sample_num = 0.0
|
|
208
|
+
for batch_data in tqdm(dataloader):
|
|
209
|
+
real = batch_data['real']
|
|
210
|
+
real_features = self._forward_modality(real, real_flag)
|
|
211
|
+
fake = batch_data['fake']
|
|
212
|
+
fake_features = self._forward_modality(fake, fake_flag)
|
|
213
|
+
|
|
214
|
+
# normalize features
|
|
215
|
+
real_features = real_features / real_features.norm(dim=1, keepdim=True).to(
|
|
216
|
+
torch.float32
|
|
217
|
+
)
|
|
218
|
+
fake_features = fake_features / fake_features.norm(dim=1, keepdim=True).to(
|
|
219
|
+
torch.float32
|
|
220
|
+
)
|
|
221
|
+
|
|
222
|
+
# calculate scores
|
|
223
|
+
score = (fake_features * real_features).sum()
|
|
224
|
+
score_acc += score
|
|
225
|
+
sample_num += real_features.shape[0]
|
|
226
|
+
|
|
227
|
+
clip_score = score_acc / sample_num
|
|
228
|
+
return clip_score.cpu().item()
|
|
229
|
+
|
|
230
|
+
def _evaluate_single_file(
|
|
231
|
+
self, image_path, text_path_or_string, image_flag='img', text_flag='txt'
|
|
232
|
+
):
|
|
233
|
+
"""Evaluate CLIP score between a single image file and text
|
|
234
|
+
|
|
235
|
+
Args:
|
|
236
|
+
image_path: Path to image file
|
|
237
|
+
text_path_or_string: Path to text file or text string directly
|
|
238
|
+
image_flag: Type of image input modality, default is 'img'
|
|
239
|
+
text_flag: Type of text input modality, default is 'txt'
|
|
240
|
+
|
|
241
|
+
Returns:
|
|
242
|
+
float: CLIP score
|
|
243
|
+
"""
|
|
244
|
+
# Determine which is image and which is text
|
|
245
|
+
if image_flag == 'img' and text_flag == 'txt':
|
|
246
|
+
img_path, txt_path = image_path, text_path_or_string
|
|
247
|
+
img_flag, txt_flag = image_flag, text_flag
|
|
248
|
+
elif image_flag == 'txt' and text_flag == 'img':
|
|
249
|
+
img_path, txt_path = text_path_or_string, image_path
|
|
250
|
+
img_flag, txt_flag = text_flag, image_flag
|
|
251
|
+
else:
|
|
252
|
+
raise ValueError("Must specify one 'img' and one 'txt' modality")
|
|
253
|
+
|
|
254
|
+
# Create a single sample dataset
|
|
255
|
+
dataset = DummyDataset(
|
|
256
|
+
img_path,
|
|
257
|
+
txt_path,
|
|
258
|
+
img_flag,
|
|
259
|
+
txt_flag,
|
|
260
|
+
transform=self.processor,
|
|
261
|
+
tokenizer=self.tokenizer,
|
|
262
|
+
)
|
|
263
|
+
|
|
264
|
+
# Get data
|
|
265
|
+
sample = dataset[0]
|
|
266
|
+
img_data = sample['real'] if img_flag == 'real_flag' else sample['fake']
|
|
267
|
+
txt_data = sample['fake'] if txt_flag == 'txt' else sample['real']
|
|
268
|
+
|
|
269
|
+
# Compute features
|
|
270
|
+
img_features = self._forward_modality(img_data, 'img')
|
|
271
|
+
txt_features = self._forward_modality(txt_data, 'txt')
|
|
272
|
+
|
|
273
|
+
# Normalize features
|
|
274
|
+
img_features = img_features / img_features.norm(dim=1, keepdim=True).to(
|
|
275
|
+
torch.float32
|
|
276
|
+
)
|
|
277
|
+
txt_features = txt_features / txt_features.norm(dim=1, keepdim=True).to(
|
|
278
|
+
torch.float32
|
|
279
|
+
)
|
|
280
|
+
|
|
281
|
+
# Compute score
|
|
282
|
+
score = (img_features * txt_features).sum()
|
|
283
|
+
|
|
284
|
+
return score.cpu().item()
|
|
285
|
+
|
|
286
|
+
def _forward_modality(self, data, flag):
|
|
287
|
+
device = self.device
|
|
288
|
+
for key in data:
|
|
289
|
+
data[key] = data[key].to(device)
|
|
290
|
+
if flag == 'img':
|
|
291
|
+
features = self.model.get_image_features(**data)
|
|
292
|
+
elif flag == 'txt':
|
|
293
|
+
features = self.model.get_text_features(**data)
|
|
294
|
+
else:
|
|
295
|
+
raise TypeError(f'Got unexpected modality: {flag}')
|
|
296
|
+
return features
|
|
@@ -0,0 +1,100 @@
|
|
|
1
|
+
# Taken from https://github.com/LAION-AI/aesthetic-predictor/blob/main/asthetics_predictor.ipynb
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
import torch
|
|
5
|
+
import open_clip
|
|
6
|
+
import torch.nn as nn
|
|
7
|
+
from PIL import Image
|
|
8
|
+
from os.path import expanduser # pylint: disable=import-outside-toplevel
|
|
9
|
+
from urllib.request import urlretrieve # pylint: disable=import-outside-toplevel
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class LaionAIAestheticPredictor:
|
|
13
|
+
def __init__(self, model_name="vit_l_14"):
|
|
14
|
+
"""Initialize the aesthetic predictor with a specified model."""
|
|
15
|
+
self.model_name = model_name
|
|
16
|
+
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
17
|
+
self.aes_model = self.get_aesthetic_model()
|
|
18
|
+
|
|
19
|
+
def get_aesthetic_model(self):
|
|
20
|
+
"""Load the aesthetic model based on the model type defined in __init__."""
|
|
21
|
+
home = expanduser("~")
|
|
22
|
+
cache_folder = home + "/.cache/emb_reader"
|
|
23
|
+
path_to_model = cache_folder + f"/sa_0_4_{self.model_name}_linear.pth"
|
|
24
|
+
|
|
25
|
+
if not os.path.exists(path_to_model):
|
|
26
|
+
os.makedirs(cache_folder, exist_ok=True)
|
|
27
|
+
url_model = f"https://github.com/LAION-AI/aesthetic-predictor/blob/main/sa_0_4_{self.model_name}_linear.pth?raw=true"
|
|
28
|
+
urlretrieve(url_model, path_to_model)
|
|
29
|
+
|
|
30
|
+
if self.model_name == "vit_l_14":
|
|
31
|
+
m = nn.Linear(768, 1)
|
|
32
|
+
elif self.model_name == "vit_b_32":
|
|
33
|
+
m = nn.Linear(512, 1)
|
|
34
|
+
else:
|
|
35
|
+
raise ValueError(f"Unsupported model: {self.model_name}")
|
|
36
|
+
|
|
37
|
+
s = torch.load(path_to_model, map_location=self.device)
|
|
38
|
+
m.load_state_dict(s)
|
|
39
|
+
m.to(self.device)
|
|
40
|
+
m.eval()
|
|
41
|
+
return m
|
|
42
|
+
|
|
43
|
+
def evaluate_aesthetic_score(self, image_path):
|
|
44
|
+
"""Evaluate the aesthetic score of a single image
|
|
45
|
+
|
|
46
|
+
Args:
|
|
47
|
+
image_path: Path to the image file
|
|
48
|
+
|
|
49
|
+
Returns:
|
|
50
|
+
float: Aesthetic score of the image
|
|
51
|
+
"""
|
|
52
|
+
if self.aes_model is None:
|
|
53
|
+
self.aes_model = self.get_aesthetic_model()
|
|
54
|
+
|
|
55
|
+
model, _, preprocess = open_clip.create_model_and_transforms(
|
|
56
|
+
'ViT-L-14', pretrained='openai'
|
|
57
|
+
)
|
|
58
|
+
try:
|
|
59
|
+
image = Image.open(image_path).convert('RGB')
|
|
60
|
+
image_tensor = preprocess(image).unsqueeze(0).to(self.device)
|
|
61
|
+
|
|
62
|
+
with torch.no_grad():
|
|
63
|
+
image_features = model.encode_image(image_tensor)
|
|
64
|
+
image_features /= image_features.norm(dim=-1, keepdim=True)
|
|
65
|
+
score = self.aes_model(image_features)
|
|
66
|
+
|
|
67
|
+
return score[0][0].item()
|
|
68
|
+
except Exception as e:
|
|
69
|
+
print(f"Error evaluating image {image_path}: {e}")
|
|
70
|
+
return None
|
|
71
|
+
|
|
72
|
+
def evaluate_folder_aesthetic_score(self, folder_path):
|
|
73
|
+
"""Evaluate the average aesthetic score of all images in a folder
|
|
74
|
+
|
|
75
|
+
Args:
|
|
76
|
+
folder_path: Path to folder containing images
|
|
77
|
+
|
|
78
|
+
Returns:
|
|
79
|
+
float: Average aesthetic score of all images in the folder
|
|
80
|
+
"""
|
|
81
|
+
if not os.path.isdir(folder_path):
|
|
82
|
+
raise ValueError(f"{folder_path} is not a valid folder path")
|
|
83
|
+
|
|
84
|
+
image_files = [
|
|
85
|
+
f
|
|
86
|
+
for f in os.listdir(folder_path)
|
|
87
|
+
if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.webp'))
|
|
88
|
+
]
|
|
89
|
+
|
|
90
|
+
if not image_files:
|
|
91
|
+
return None
|
|
92
|
+
|
|
93
|
+
scores = []
|
|
94
|
+
for img_file in image_files:
|
|
95
|
+
img_path = os.path.join(folder_path, img_file)
|
|
96
|
+
score = self.evaluate_aesthetic_score(img_path)
|
|
97
|
+
if score is not None:
|
|
98
|
+
scores.append(score)
|
|
99
|
+
|
|
100
|
+
return sum(scores) / len(scores) if scores else None
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
import argparse
|
|
2
|
+
import os
|
|
3
|
+
from image_evaluator.laion_ai_aesthetic_predictor import LaionAIAestheticPredictor
|
|
4
|
+
from image_evaluator.clip_score_predictor import ClipScorePredictor
|
|
5
|
+
from image_evaluator.arcface_dist_predictor import ArcFaceDistPredictor
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def main():
|
|
9
|
+
parser = argparse.ArgumentParser(
|
|
10
|
+
description="Evaluate the aesthetic score of an image."
|
|
11
|
+
)
|
|
12
|
+
parser.add_argument("--image", type=str, help="Path to the image file or folder")
|
|
13
|
+
parser.add_argument(
|
|
14
|
+
"--prompt",
|
|
15
|
+
type=str,
|
|
16
|
+
default=None,
|
|
17
|
+
help="Path to the prompt file or text prompt",
|
|
18
|
+
)
|
|
19
|
+
parser.add_argument(
|
|
20
|
+
"--reference",
|
|
21
|
+
type=str,
|
|
22
|
+
default=None,
|
|
23
|
+
help="Path to the reference image for ArcFace distance",
|
|
24
|
+
)
|
|
25
|
+
args = parser.parse_args()
|
|
26
|
+
|
|
27
|
+
# Check if it's a file or folder
|
|
28
|
+
is_folder = os.path.isdir(args.image) if args.image else False
|
|
29
|
+
|
|
30
|
+
# LAION AI Aesthetic Score
|
|
31
|
+
laion_ai_aesthetic_predictor = LaionAIAestheticPredictor()
|
|
32
|
+
if is_folder:
|
|
33
|
+
laion_ai_aesthetic_score = (
|
|
34
|
+
laion_ai_aesthetic_predictor.evaluate_folder_aesthetic_score(args.image)
|
|
35
|
+
)
|
|
36
|
+
else:
|
|
37
|
+
laion_ai_aesthetic_score = (
|
|
38
|
+
laion_ai_aesthetic_predictor.evaluate_aesthetic_score(args.image)
|
|
39
|
+
)
|
|
40
|
+
|
|
41
|
+
# CLIP Score Evaluation
|
|
42
|
+
clip_score_predictor = ClipScorePredictor()
|
|
43
|
+
clip_score = clip_score_predictor.evaluate_clip_score(args.image, args.prompt)
|
|
44
|
+
|
|
45
|
+
# ArcFace Distance Evaluation
|
|
46
|
+
arcface_distance_predictor = ArcFaceDistPredictor()
|
|
47
|
+
if is_folder and args.reference and os.path.isdir(args.reference):
|
|
48
|
+
arcface_distance = arcface_distance_predictor.evaluate_folder_arcface_distance(
|
|
49
|
+
args.reference, args.image
|
|
50
|
+
)
|
|
51
|
+
elif args.reference:
|
|
52
|
+
arcface_distance = arcface_distance_predictor.evaluate_arcface_distance(
|
|
53
|
+
args.reference, args.image
|
|
54
|
+
)
|
|
55
|
+
else:
|
|
56
|
+
arcface_distance = None
|
|
57
|
+
|
|
58
|
+
print(f"LAION AI Aesthetic Score: {laion_ai_aesthetic_score}")
|
|
59
|
+
print(f"CLIP Score: {clip_score}")
|
|
60
|
+
if arcface_distance is not None:
|
|
61
|
+
print(f"ArcFace Distance: {arcface_distance}")
|
|
62
|
+
else:
|
|
63
|
+
print("ArcFace Distance: Not evaluated (reference image required)")
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
if __name__ == "__main__":
|
|
67
|
+
main()
|
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
[tool.poetry]
|
|
2
|
+
name = "image-evaluator"
|
|
3
|
+
version = "0.0.1"
|
|
4
|
+
description = "An automatic image evaluation script"
|
|
5
|
+
authors = ["neverbiasu"]
|
|
6
|
+
license = "MIT"
|
|
7
|
+
readme = "README.md"
|
|
8
|
+
packages = [{include = "image_evaluator"}]
|
|
9
|
+
|
|
10
|
+
[tool.poetry.dependencies]
|
|
11
|
+
python = "^3.8"
|
|
12
|
+
clip-score = "*"
|
|
13
|
+
insightface = "*"
|
|
14
|
+
open-clip-torch = "*"
|
|
15
|
+
torch = ">=1.7.0"
|
|
16
|
+
transformers = "*"
|
|
17
|
+
pillow = "*"
|
|
18
|
+
numpy = "*"
|
|
19
|
+
tqdm = "*"
|
|
20
|
+
|
|
21
|
+
[tool.poetry.group.dev.dependencies]
|
|
22
|
+
pytest = "^7.0.0"
|
|
23
|
+
black = "^23.0.0"
|
|
24
|
+
isort = "^5.12.0"
|
|
25
|
+
|
|
26
|
+
[build-system]
|
|
27
|
+
requires = ["poetry-core"]
|
|
28
|
+
build-backend = "poetry.core.masonry.api"
|
|
29
|
+
|
|
30
|
+
[tool.poetry.scripts]
|
|
31
|
+
image-evaluator = "image_evaluator.main:main"
|