Update app.py

app.py

@@ -10,6 +10,7 @@ from imgutils.utils import open_onnx_model
 
 import random
 from typing import List
+from dataclasses import dataclass
 from base import ONNXBaseTask
 from utils import prepare_input_wraper
 
@@ -29,6 +30,40 @@ def _onnx_model(name):
         f'{name}'
     ))
 
+@dataclass
+class NSFWPrediction:
+    label: str
+    score: float
+
+class CheckNSFWTask(ONNXBaseTask):
+    classes: List[str] = ['drawings', 'hentai', 'neutral', 'porn', 'sexy']
+
+    def __init__(self, weight: str):
+        super().__init__(weight)
+
+    def __call__(self, image) -> NSFWPrediction:
+        return self.call(image)
+
+    def process_output(self, raw_outputs) -> NSFWPrediction:
+        probabilities = raw_outputs[0][0]
+        max_prob_index = np.argmax(probabilities)
+        max_prob_score = probabilities[max_prob_index]
+
+        predicted_class = self.classes[max_prob_index]
+
+        return NSFWPrediction(label=predicted_class, score=max_prob_score.tolist())
+
+    def setup_prepare_input(self):
+        return prepare_input_wraper(
+            inter=1,
+            color_space="RGB",
+            mean=None,
+            std=None,
+            is_scale=True,
+            channel_first=False
+        )
+
+
 def _image_preprocess(image, size: int = 224) -> np.ndarray:
     image = load_image(image, mode='RGB').resize((size, size), Image.NEAREST)
     return (np.array(image) / 255.0)[None, ...]
@@ -55,7 +90,7 @@ if __name__ == '__main__':
                 gr_ratings = gr.Label(label='Ratings')
 
         gr_btn_submit.click(
-            predict,
+            CheckNSFWTask,
             inputs=[gr_input_image, gr_model],
             outputs=[gr_ratings],
         )