Upload 2 files

app.py

@@ -1,31 +1,57 @@
-import gradio as gr
-import torch
-from transformers import AutoImageProcessor, ConvNextV2ForImageClassification
-from transformers import AutoModelForImageClassification
-from torch import nn
-import dbimutils as utils
-
-DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-image_processor  = AutoImageProcessor.from_pretrained("Muinez/artwork-scorer")
-model = AutoModelForImageClassification.from_pretrained("Muinez/artwork-scorer", problem_type="multi_label_classification").to(DEVICE)
-
-def predict(img):
-	file = utils.preprocess_image(img)
-	encoded = image_processor(file, return_tensors="pt").to(DEVICE)
-
-	with torch.no_grad():
-		logits = model(**encoded).logits.cpu()
-    
-	outputs = nn.functional.sigmoid(logits)
-
-	return outputs[0][0].item(), outputs[0][1].item(), outputs[0][2].item()
-
-gr.Interface(
-    title="Artwork scorer",
-    description="Predicts score (0-1) for artwork.\nCould be wrong!!!\nDoes not work very well with nsfw i.e. it was not trained on it",
-    fn=predict,
-	allow_flagging="never",
-    inputs=gr.Image(type="pil"),
-    outputs=[gr.Number(label="Score"), gr.Number(label="View count ratio (probably useless)"), gr.Number(label="Upload date 0 - 2016, 1 - 2023")]
+import gradio as gr
+import torch
+from torch import nn
+from transformers import SiglipImageProcessor,SiglipModel
+import dbimutils as utils
+
+class ScoreClassifier(nn.Module):
+	def __init__(self):
+		super(ScoreClassifier, self).__init__()
+
+		self.classifier = nn.Sequential(
+			nn.Linear(256, 1),
+			nn.Sigmoid()
+		)
+
+		self.extractor = nn.Sequential(
+            nn.Linear(768, 512),
+            nn.BatchNorm1d(512),
+			nn.ReLU(),
+            nn.Linear(512, 256),
+            nn.BatchNorm1d(256),
+			nn.ReLU(),
+            nn.Linear(256, 256),
+			nn.ReLU(),
+            )
+
+	def forward(self, img):
+		return self.classifier(self.extractor(img))
+
+from huggingface_hub import hf_hub_download
+model_file = hf_hub_download(repo_id="Muinez/Image-scorer", filename="scorer.pth")
+
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+model = ScoreClassifier().to(DEVICE)
+model.load_state_dict(torch.load("scorer.pth"))
+model.eval()
+
+processor = SiglipImageProcessor.from_pretrained('google/siglip-base-patch16-512')
+siglip = SiglipModel.from_pretrained('google/siglip-base-patch16-512').to(DEVICE)
+
+def predict(img):
+	img = utils.preprocess_image(img)
+	encoded = processor(img, return_tensors="pt").pixel_values.to(DEVICE)
+
+	with torch.no_grad():
+		score = model(siglip.get_image_features(encoded))
+    
+	return score.item()
+
+gr.Interface(
+    title="Artwork scorer",
+    description="Predicts score (0-1) for artwork.\nCould be wrong!!!\nDoes not work very well with nsfw i.e. it was not trained on it",
+    fn=predict,
+	allow_flagging="never",
+    inputs=gr.Image(type="pil"),
+    outputs=[gr.Number(label="Score")]
 ).launch()

dbimutils.py

@@ -61,8 +61,8 @@ def preprocess_image(img):
     image = new_image.convert('RGB')
     image = np.asarray(image)
 
-    image = make_square(image, 384)
-    image = smart_resize(image, 384)
+    image = make_square(image, 512)
+    image = smart_resize(image, 512)
     image = image.astype(np.float32)
 
     return Image.fromarray(np.uint8(image))