Upload 2 files

app.py

@@ -0,0 +1,33 @@
+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], outputs[0][1]
+
+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)")]
+).launch()
+
+

dbimutils.py

@@ -0,0 +1,68 @@
+# DanBooru IMage Utility functions
+# Taken from https://huggingface.co/spaces/SmilingWolf/wd-v1-4-tags
+
+import cv2
+import numpy as np
+from PIL import Image
+import PIL
+
+def smart_imread(img, flag=cv2.IMREAD_UNCHANGED):
+    if img.endswith(".gif"):
+        img = Image.open(img)
+        img = img.convert("RGB")
+        img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+    else:
+        img = cv2.imread(img, flag)
+    return img
+
+
+def smart_24bit(img):
+    if img.dtype is np.dtype(np.uint16):
+        img = (img / 257).astype(np.uint8)
+
+    if len(img.shape) == 2:
+        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+    elif img.shape[2] == 4:
+        trans_mask = img[:, :, 3] == 0
+        img[trans_mask] = [255, 255, 255, 255]
+        img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
+    return img
+
+
+def make_square(img, target_size):
+    old_size = img.shape[:2]
+    desired_size = max(old_size)
+    desired_size = max(desired_size, target_size)
+
+    delta_w = desired_size - old_size[1]
+    delta_h = desired_size - old_size[0]
+    top, bottom = delta_h // 2, delta_h - (delta_h // 2)
+    left, right = delta_w // 2, delta_w - (delta_w // 2)
+
+    color = [255, 255, 255]
+    new_im = cv2.copyMakeBorder(
+        img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color
+    )
+    return new_im
+
+
+def smart_resize(img, size):
+    # Assumes the image has already gone through make_square
+    if img.shape[0] > size:
+        img = cv2.resize(img, (size, size), interpolation=cv2.INTER_AREA)
+    elif img.shape[0] < size:
+        img = cv2.resize(img, (size, size), interpolation=cv2.INTER_CUBIC)
+    return img
+
+def preprocess_image(img):
+    image = img.convert('RGBA')
+    new_image = PIL.Image.new('RGBA', image.size, 'WHITE')
+    new_image.paste(image, mask=image)
+    image = new_image.convert('RGB')
+    image = np.asarray(image)
+
+    image = make_square(image, 384)
+    image = smart_resize(image, 384)
+    image = image.astype(np.float32)
+
+    return Image.fromarray(np.uint8(image))