File size: 2,348 Bytes
6d6dbbc | 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 | """
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)
# normalize embedding
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]))
|