|
import os |
|
import numpy as np |
|
import torch |
|
import pytorch_lightning as pl |
|
import torch.nn as nn |
|
import clip |
|
from PIL import Image, ImageFile |
|
import gradio as gr |
|
|
|
|
|
class MLP(pl.LightningModule): |
|
def __init__(self, input_size, xcol='emb', ycol='avg_rating'): |
|
super().__init__() |
|
self.input_size = input_size |
|
self.xcol = xcol |
|
self.ycol = ycol |
|
self.layers = nn.Sequential( |
|
nn.Linear(self.input_size, 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) |
|
) |
|
|
|
def forward(self, x): |
|
return self.layers(x) |
|
|
|
def training_step(self, batch, batch_idx): |
|
x = batch[self.xcol] |
|
y = batch[self.ycol].reshape(-1, 1) |
|
x_hat = self.layers(x) |
|
loss = F.mse_loss(x_hat, y) |
|
return loss |
|
|
|
def validation_step(self, batch, batch_idx): |
|
x = batch[self.xcol] |
|
y = batch[self.ycol].reshape(-1, 1) |
|
x_hat = self.layers(x) |
|
loss = F.mse_loss(x_hat, y) |
|
return loss |
|
|
|
def configure_optimizers(self): |
|
optimizer = torch.optim.Adam(self.parameters(), lr=1e-3) |
|
return optimizer |
|
|
|
def normalized(a, axis=-1, order=2): |
|
import numpy as np |
|
|
|
l2 = np.atleast_1d(np.linalg.norm(a, order, axis)) |
|
l2[l2 == 0] = 1 |
|
return a / np.expand_dims(l2, axis) |
|
|
|
def load_models(): |
|
model = MLP(768) |
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
s = torch.load("sac+logos+ava1-l14-linearMSE.pth", map_location=device) |
|
|
|
model.load_state_dict(s) |
|
model.to(device) |
|
model.eval() |
|
|
|
model2, preprocess = clip.load("ViT-L/14", device=device) |
|
|
|
model_dict = {} |
|
model_dict['classifier'] = model |
|
model_dict['clip_model'] = model2 |
|
model_dict['clip_preprocess'] = preprocess |
|
model_dict['device'] = device |
|
|
|
return model_dict |
|
|
|
def predict(image): |
|
image_input = model_dict['clip_preprocess'](image).unsqueeze(0).to(model_dict['device']) |
|
with torch.no_grad(): |
|
image_features = model_dict['clip_model'].encode_image(image_input) |
|
if model_dict['device'] == 'cuda': |
|
im_emb_arr = normalized(image_features.detach().cpu().numpy()) |
|
im_emb = torch.from_numpy(im_emb_arr).to(model_dict['device']).type(torch.cuda.FloatTensor) |
|
else: |
|
im_emb_arr = normalized(image_features.detach().numpy()) |
|
im_emb = torch.from_numpy(im_emb_arr).to(model_dict['device']).type(torch.FloatTensor) |
|
|
|
prediction = model_dict['classifier'](im_emb) |
|
score = prediction.item() |
|
|
|
return {'aesthetic score': score} |
|
|
|
if __name__ == '__main__': |
|
print('\tinit models') |
|
|
|
global model_dict |
|
|
|
model_dict = load_models() |
|
|
|
inputs = [gr.inputs.Image(type='pil', label='Image')] |
|
|
|
outputs = gr.outputs.JSON() |
|
|
|
title = 'image aesthetic predictor' |
|
|
|
examples = ['example1.jpg', 'example2.jpg', 'example3.jpg'] |
|
|
|
description = """ |
|
# Image Aesthetic Predictor Demo |
|
This model (Image Aesthetic Predictor) is trained by LAION Team. See [https://github.com/christophschuhmann/improved-aesthetic-predictor](https://github.com/christophschuhmann/improved-aesthetic-predictor) |
|
1. This model is desgined by adding five MLP layers on top of (frozen) CLIP ViT-L/14 and only the MLP layers are fine-tuned with a lot of images by a regression loss term such as MSE and MAE. |
|
2. Output is bounded from 0 to 10. The higher the better. |
|
""" |
|
|
|
article = "<p style='text-align: center'><a href='https://laion.ai/blog/laion-aesthetics/'>LAION aesthetics blog post</a></p>" |
|
|
|
with gr.Blocks() as demo: |
|
gr.Markdown(description) |
|
with gr.Row(): |
|
with gr.Column(): |
|
image_input = gr.Image(type='pil', label='Input image') |
|
submit_button = gr.Button('Submit') |
|
json_output = gr.JSON(label='Output') |
|
submit_button.click(predict, inputs=image_input, outputs=json_output) |
|
gr.Examples(examples=examples, inputs=image_input) |
|
gr.HTML(article) |
|
demo.launch() |
|
|
