Update app.py

app.py

@@ -1,3 +1,4 @@
+
 import os
 from functools import lru_cache
 
@@ -8,13 +9,6 @@ from huggingface_hub import hf_hub_download
 from imgutils.data import load_image
 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
-
-
 _MODELS = [
     ('content_moderation.onnx', 224),
 ]
@@ -30,39 +24,6 @@ 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)
@@ -74,7 +35,7 @@ _LABELS = ['drawings', 'hentai', 'neutral', 'porn', 'sexy']
 
 def predict(image, model_name):
     input_ = _image_preprocess(image, _MODEL_TO_SIZE[model_name]).astype(np.float32)
-    output_, = _onnx_model(model_name).run(['dense_3'], {'input_1': input_})
+    output_, = _onnx_model(model_name).run()
     return dict(zip(_LABELS, map(float, output_[0])))
 
 
@@ -90,7 +51,7 @@ if __name__ == '__main__':
                 gr_ratings = gr.Label(label='Ratings')
 
         gr_btn_submit.click(
-            CheckNSFWTask,
+            predict,
             inputs=[gr_input_image, gr_model],
             outputs=[gr_ratings],
         )