| """ |
| Adapted from https://github.com/christophschuhmann/improved-aesthetic-predictor. Originally Apache License, Version 2.0, January 2004. |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| import numpy as np |
| from transformers import CLIPModel, CLIPProcessor |
| from PIL import Image |
| from io import BytesIO |
|
|
| def open_image(image): |
| if isinstance(image, bytes): |
| image = Image.open(BytesIO(image)) |
| elif isinstance(image, str): |
| image = Image.open(image) |
| image = image.convert("RGB") |
| return image |
|
|
|
|
| class MLP(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.layers = nn.Sequential( |
| nn.Linear(768, 1024), |
| nn.Dropout(0.2), |
| nn.Linear(1024, 128), |
| nn.Dropout(0.2), |
| nn.Linear(128, 64), |
| nn.Dropout(0.1), |
| nn.Linear(64, 16), |
| nn.Linear(16, 1), |
| ) |
|
|
| @torch.no_grad() |
| def forward(self, embed): |
| return self.layers(embed) |
|
|
|
|
| class AestheticScorer(torch.nn.Module): |
| def __init__(self, dtype, clip_name_or_path, aesthetic_name_or_path): |
| super().__init__() |
| self.clip = CLIPModel.from_pretrained(clip_name_or_path) |
| self.processor = CLIPProcessor.from_pretrained(clip_name_or_path) |
| self.mlp = MLP() |
| state_dict = torch.load(aesthetic_name_or_path, map_location='cpu') |
| self.mlp.load_state_dict(state_dict) |
| self.dtype = dtype |
| self.eval() |
|
|
| @torch.no_grad() |
| def __call__(self, images): |
| device = next(self.parameters()).device |
| inputs = self.processor(images=images, return_tensors="pt") |
| inputs = {k: v.to(self.dtype).to(device) for k, v in inputs.items()} |
| embed = self.clip.get_image_features(**inputs) |
| |
| embed = embed / torch.linalg.vector_norm(embed, dim=-1, keepdim=True) |
| return self.mlp(embed).squeeze(1) |
| |
| |
| if __name__ == "__main__": |
| clip_name_or_path = "openai/clip-vit-large-patch14" |
| aesthetic_name_or_path = "./sac+logos+ava1-l14-linearMSE.pth" |
| aesthetic_scorer = AestheticScorer(torch.float32, clip_name_or_path, aesthetic_name_or_path) |
| image0 = open_image('./image0.png') |
| image1 = open_image('./image1.png') |
| print(aesthetic_scorer(image0)) |
| print(aesthetic_scorer(image1)) |
| print(aesthetic_scorer([image0, image1])) |
| |
