Update app.py

app.py

@@ -1,17 +1,33 @@
 from transformers import pipeline
+from imgutils.data import rgb_encode, load_image
+from onnx_ import _open_onnx_model
+from PIL import Image
 import gradio as gr
+import numpy as np
 import os
 import requests
 import timm
 import torch
 import json
 
+def _img_encode(image, size=(384,384), normalize=(0.5,0.5)):
+    image = image.resize(size, Image.BILINEAR)
+    data = rgb_encode(image, order_='CHW')
+
+    if normalize is not None:
+        mean_, std_ = normalize
+        mean = np.asarray([mean_]).reshape((-1, 1, 1))
+        std = np.asarray([std_]).reshape((-1, 1, 1))
+        data = (data - mean) / std
+
+    return data.astype(np.float32)
+    
 nsfw_tf = pipeline(model="carbon225/vit-base-patch16-224-hentai")
 
-if not os.path.exists("timm.ckpt"):
-    open("timm.ckpt", "wb").write(
+if not os.path.exists("timm.onnx"):
+    open("timm.onnx", "wb").write(
         requests.get(
-            "https://huggingface.co/deepghs/anime_rating/resolve/main/caformer_s36_plus/model.ckpt"
+            "https://huggingface.co/deepghs/anime_rating/resolve/main/caformer_s36_plus/model.onnx"
         ).content
     )
     open("timmcfg.json", "wb").write(
@@ -25,26 +41,14 @@ else:
 with open("timmcfg.json") as file:
     tm_cfg = json.load(file)
 
-nsfw_tm = timm.create_model(
-    "caformer_s36.sail_in22k_ft_in1k_384",
-    checkpoint_path="./timm.ckpt",
-    model_config=tm_cfg,
-    num_classes=3
-).eval()
-tm_config = timm.data.resolve_model_data_config(nsfw_tm)
-tm_trans = timm.data.create_transform((256, 256), **tm_config, is_training=False)
-
+nsfw_tm = _open_onnx_model("timm.onnx")
 
 def launch(img):
     weight = 0
-    img = Image.open(img).convert('RGB')
-    tm_output = model.pretrained_cfg['labels'][
-        torch.argmax(
-            torch.nn.functional.softmax(
-                nsfw_tm(transforms(img).unsqueeze(0))[0], dim=0
-            )
-        )
-    ]
+    tm_image = load_image(img, mode='RGB')
+    tm_input_ = _img_encode(tm_image, size=(256, 256))[None, ...]
+    tm_output, = nsfw_tm.run(['output'], {'input': tm_input_})
+    tm_output = zip(tm_cfg["labels"], map(lambda x: x.item(), output[0]))[0][0]
     
     match tm_output:
         case "safe":
@@ -54,8 +58,8 @@ def launch(img):
         case "r18":
             weight += 2
 
-
-    tf_output = nsfw_tf(img)[0]["label"]
+    tf_img = Image.open(img).convert('RGB')
+    tf_output = nsfw_tf(tf_img)[0]["label"]
 
     match tf_output:
         case "safe":