Duplicate from Geonmo/laion-aesthetic-predictor

Co-authored-by: Geonmo Gu <[email protected]>

.gitattributes

@@ -0,0 +1,33 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,14 @@
+---
+title: Laion Aesthetic Predictor
+emoji: 🎨
+colorFrom: red
+colorTo: indigo
+sdk: gradio
+sdk_version: 3.6
+app_file: app.py
+pinned: false
+license: apache-2.0
+duplicated_from: Geonmo/laion-aesthetic-predictor
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,133 @@
+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
+
+# if you changed the MLP architecture during training, change it also here:
+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.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(1024, 128),
+            #nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(128, 64),
+            #nn.ReLU(),
+            nn.Dropout(0.1),
+
+            nn.Linear(64, 16),
+            #nn.ReLU(),
+
+            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  # pylint: disable=import-outside-toplevel
+
+    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()

requirements.txt

@@ -0,0 +1,6 @@
+ftfy
+regex
+git+https://github.com/openai/CLIP.git
+gradio
+torch
+pytorch-lightning

sac+logos+ava1-l14-linearMSE.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21dd590f3ccdc646f0d53120778b296013b096a035a2718c9cb0d511bff0f1e0
+size 3714759