Duplicate from suko/nsfw

Co-authored-by: SuKo Kuo <[email protected]>

.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,14 @@
+---
+title: Nsfw Space
+emoji: 💩
+colorFrom: green
+colorTo: green
+sdk: gradio
+sdk_version: 3.21.0
+app_file: app.py
+pinned: false
+license: apache-2.0
+duplicated_from: suko/nsfw
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,105 @@
+import gradio as gr
+import numpy as np
+import json
+import os
+from PIL import Image
+import onnxruntime as rt
+class ONNXModel:
+    def __init__(self, dir_path) -> None:
+        """Method to get name of model file. Assumes model is in the parent directory for script."""
+        model_dir = os.path.dirname(dir_path)
+        with open(os.path.join(model_dir, "signature.json"), "r") as f:
+            self.signature = json.load(f)
+        self.model_file = os.path.join(model_dir, self.signature.get("filename"))
+        if not os.path.isfile(self.model_file):
+            raise FileNotFoundError(f"Model file does not exist")
+        # get the signature for model inputs and outputs
+        self.signature_inputs = self.signature.get("inputs")
+        self.signature_outputs = self.signature.get("outputs")
+        self.session = None
+        if "Image" not in self.signature_inputs:
+            raise ValueError("ONNX model doesn't have 'Image' input! Check signature.json, and please report issue to Lobe.")
+        # Look for the version in signature file.
+        # If it's not found or the doesn't match expected, print a message
+        version = self.signature.get("export_model_version")
+        if version is None or version != EXPORT_MODEL_VERSION:
+            print(
+                f"There has been a change to the model format. Please use a model with a signature 'export_model_version' that matches {EXPORT_MODEL_VERSION}."
+            )
+
+    def load(self) -> None:
+        """Load the model from path to model file"""
+        # Load ONNX model as session.
+        self.session = rt.InferenceSession(path_or_bytes=self.model_file)
+
+    def predict(self, image: Image.Image) -> dict:
+        """
+        Predict with the ONNX session!
+        """
+        # process image to be compatible with the model
+        img = self.process_image(image, self.signature_inputs.get("Image").get("shape"))
+        # run the model!
+        fetches = [(key, value.get("name")) for key, value in self.signature_outputs.items()]
+        # make the image a batch of 1
+        feed = {self.signature_inputs.get("Image").get("name"): [img]}
+        outputs = self.session.run(output_names=[name for (_, name) in fetches], input_feed=feed)
+        return self.process_output(fetches, outputs)
+
+    def process_image(self, image: Image.Image, input_shape: list) -> np.ndarray:
+        """
+        Given a PIL Image, center square crop and resize to fit the expected model input, and convert from [0,255] to [0,1] values.
+        """
+        width, height = image.size
+        # ensure image type is compatible with model and convert if not
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+        # center crop image (you can substitute any other method to make a square image, such as just resizing or padding edges with 0)
+        if width != height:
+            square_size = min(width, height)
+            left = (width - square_size) / 2
+            top = (height - square_size) / 2
+            right = (width + square_size) / 2
+            bottom = (height + square_size) / 2
+            # Crop the center of the image
+            image = image.crop((left, top, right, bottom))
+        # now the image is square, resize it to be the right shape for the model input
+        input_width, input_height = input_shape[1:3]
+        if image.width != input_width or image.height != input_height:
+            image = image.resize((input_width, input_height))
+
+        # make 0-1 float instead of 0-255 int (that PIL Image loads by default)
+        image = np.asarray(image) / 255.0
+        # format input as model expects
+        return image.astype(np.float32)
+
+    def process_output(self, fetches: dict, outputs: dict) -> dict:
+        # un-batch since we ran an image with batch size of 1,
+        # convert to normal python types with tolist(), and convert any byte strings to normal strings with .decode()
+        out_keys = ["label", "confidence"]
+        results = {}
+        for i, (key, _) in enumerate(fetches):
+            val = outputs[i].tolist()[0]
+            if isinstance(val, bytes):
+                val = val.decode()
+            results[key] = val
+        confs = results["Confidences"]
+        labels = self.signature.get("classes").get("Label")
+        output = [dict(zip(out_keys, group)) for group in zip(labels, confs)]
+        sorted_output = {"predictions": sorted(output, key=lambda k: k["confidence"], reverse=True)}
+        return sorted_output
+EXPORT_MODEL_VERSION=1
+model = ONNXModel(dir_path="model.onnx")
+model.load()
+
+def predict(image):
+    image = Image.fromarray(np.uint8(image), 'RGB')
+    prediction = model.predict(image)
+    for output in prediction["predictions"]:
+        output["confidence"] = round(output["confidence"], 4)
+    return prediction
+
+inputs = gr.inputs.Image(type="pil")
+outputs = gr.outputs.JSON()
+
+runtime=gr.Interface(title="Naked Detector",fn=predict, inputs=inputs, outputs=outputs)
+runtime.launch()

labels.txt

@@ -0,0 +1,2 @@
+Naked
+Safe

model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c68efed73e1f740c1e0847d51a37960c06c880e84acc710bb4b65ce9db6857a1
+size 99511952

requirements.txt

@@ -0,0 +1,4 @@
+gradio
+numpy
+pillow
+onnxruntime

signature.json

@@ -0,0 +1 @@
+{"doc_id": "6bfa3ac3-7170-4519-9a3d-777096d3423b", "doc_name": "NSFW", "doc_version": "c02bcb9a6cb4348ce6a87bed14754843", "format": "onnx", "version": 433, "inputs": {"Image": {"dtype": "float32", "shape": [null, 224, 224, 3], "name": "Image:0"}}, "outputs": {"Confidences": {"dtype": "float32", "shape": [null, 2], "name": "sequential/dense_2/Softmax:0"}}, "tags": [], "classes": {"Label": ["Naked", "Safe"]}, "filename": "model.onnx", "export_model_version": 1}