Upload 3 files

aesthetic_predictor_huber_ad_ep7.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4791de2a7ce900a978d66a067a60bd3924bf8dc0b85f90426dfe9673a1de68cb
+size 33252658

app.py

@@ -0,0 +1,151 @@
+import gradio as gr
+import numpy as np
+import torch
+from transformers import AutoModel, BitImageProcessor, SiglipImageProcessor, SiglipVisionModel
+from PIL import Image, ImageOps
+from sklearn.metrics.pairwise import cosine_similarity
+import torch.nn as nn
+
+device = torch.device('cpu')
+torch.set_num_threads(4)
+
+processor_d = BitImageProcessor(do_center_crop=False, do_convert_rgb=False, do_normalize=True, do_rescale=True, do_resize=False, image_mean=[0.485, 0.456, 0.406], image_std=[0.229, 0.224, 0.225], resample=3, rescale_factor=0.00392156862745098)
+model_d = AutoModel.from_pretrained('facebook/dinov2-base', attn_implementation="sdpa").to(device)
+processor_s = SiglipImageProcessor.from_pretrained('google/siglip-so400m-patch14-384')
+model_s = SiglipVisionModel.from_pretrained('google/siglip-so400m-patch14-384', attn_implementation="sdpa").to(device)
+
+class ResidualBlock(nn.Module):
+    def __init__(self, input_size):
+        super(ResidualBlock, self).__init__()
+        self.linear1 = nn.Linear(input_size, input_size // 2)
+        self.LayerNorm1 = nn.LayerNorm(input_size // 2)
+        self.activation1 = nn.Mish()
+
+        self.linear2 = nn.Linear(input_size // 2, input_size // 4)
+        self.LayerNorm2 = nn.LayerNorm(input_size // 4)
+        self.activation2 = nn.Mish()
+
+        self.linear3 = nn.Linear(input_size // 4, input_size // 2)
+        self.LayerNorm3 = nn.LayerNorm(input_size // 2)
+        self.activation3 = nn.Mish()
+
+        self.linear4 = nn.Linear(input_size // 2, input_size)
+        self.LayerNorm4 = nn.LayerNorm(input_size)
+        self.activation4 = nn.Mish()
+
+        self.shortcut = nn.Linear(input_size, input_size)
+
+    def forward(self, x):
+        identity = self.shortcut(x)
+        out = self.linear1(x)
+        out = self.LayerNorm1(out)
+        out = self.activation1(out)
+
+        out = self.linear2(out)
+        out = self.LayerNorm2(out)
+        out = self.activation2(out)
+        
+        out = self.linear3(out)
+        out = self.LayerNorm3(out)
+        out = self.activation3(out)
+        
+        out = self.linear4(out)
+        out = self.LayerNorm4(out)
+        out = self.activation4(out)
+        
+        out += identity
+
+        return out
+
+class MLP(nn.Module):
+    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(
+            ResidualBlock(self.input_size),
+            nn.Mish(),
+            nn.Linear(1920, 1)
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+mlp = MLP(1920) 
+
+s = torch.load("./aesthetic_predictor_huber_ad_ep7.pth", map_location=torch.device('cpu'))   
+
+mlp.load_state_dict(s)
+
+mlp.to(device)
+
+mlp.eval()
+
+def normalized(a, axis=-1, order=2):
+    l2 = np.atleast_1d(np.linalg.norm(a, order, axis))
+    l2[l2 == 0] = 1
+    return a / np.expand_dims(l2, axis)
+
+def process_image(image, device):
+    image = image.convert('RGBA')
+    background = Image.new('RGBA', image.size, (255, 255, 255, 255))
+    image = Image.alpha_composite(background, image).convert('RGB')
+
+    max_side = 518
+    ratio = max_side / max(image.size)
+    new_size = (int(image.size[0] * ratio), int(image.size[1] * ratio))
+    image_d = image.resize(new_size, Image.LANCZOS)
+    max_side_s = 384
+    ratio_s = max_side_s / max(image.size)
+    new_size_s = (int(image.size[0] * ratio_s), int(image.size[1] * ratio_s))
+    image_resized = image.resize(new_size_s, Image.LANCZOS)
+
+    image_s = ImageOps.pad(image_resized, (384, 384), color=(255, 255, 255))
+
+    inputs_d = processor_d(image_d, return_tensors="pt").to(device)
+    inputs_s = processor_s(image_s, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+        outputs_d = model_d(**inputs_d)
+        outputs_s = model_s(**inputs_s)
+    class_token_d = normalized(outputs_d.pooler_output.cpu().detach().numpy())
+    class_token_s = normalized(outputs_s.pooler_output.cpu().detach().numpy())
+    im_emb_arr = np.concatenate((class_token_s, class_token_d), axis=1)
+
+    prediction_value = mlp(torch.from_numpy(im_emb_arr).to(device).type(torch.FloatTensor)).item()
+
+    return im_emb_arr, prediction_value
+
+
+def infer(image1, image2):
+    try:
+
+        features1, prediction_value1 = process_image(image1, device)
+        features2, prediction_value2 = process_image(image2, device)
+
+        cos_sim_features = cosine_similarity(features1, features2)[0][0]
+
+        return cos_sim_features, prediction_value1, prediction_value2
+    except Exception as e:
+        print(f"Error during inference: {e}")
+        return "Error", "Error", "Error"
+
+
+with gr.Blocks() as iface:
+    gr.Markdown("# Anime Aesthetic Predictor Based on Twitter User Preferences\nUpload two images to calculate the aesthetic score (0-10).")
+    with gr.Row():
+        image1 = gr.Image(type="pil")
+        image2 = gr.Image(type="pil")
+    with gr.Row():
+        prediction1 = gr.Textbox(label="Aesthetic Score 1")
+        prediction2 = gr.Textbox(label="Aesthetic Score 2")
+    with gr.Row():
+        feature_similarity = gr.Textbox(label="Feature Similarity")
+    with gr.Row():
+        submit_btn = gr.Button("Submit")
+
+    submit_btn.click(infer, inputs=[image1, image2], outputs=[feature_similarity, prediction1, prediction2])
+    
+iface.queue(max_size=10)
+iface.launch()

requirements.txt

@@ -0,0 +1,6 @@
+gradio
+numpy<3
+torch
+transformers
+Pillow
+scikit-learn