Spaces:
Running
Running
| 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: | |
| """Load model metadata and initialize ONNX session.""" | |
| 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("Model file does not exist.") | |
| self.signature_inputs = self.signature.get("inputs") | |
| self.signature_outputs = self.signature.get("outputs") | |
| if "Image" not in self.signature_inputs: | |
| raise ValueError("ONNX model must have an 'Image' input. Check signature.json.") | |
| # Check export version | |
| version = self.signature.get("export_model_version") | |
| if version is None or version != EXPORT_MODEL_VERSION: | |
| print(f"Warning: Expected model version {EXPORT_MODEL_VERSION}, but found {version}.") | |
| self.session = None | |
| def load(self) -> None: | |
| """Load the ONNX model with execution providers.""" | |
| self.session = rt.InferenceSession(self.model_file, providers=["CPUExecutionProvider"]) | |
| def predict(self, image: Image.Image) -> dict: | |
| """Process image and run ONNX model inference.""" | |
| img = self.process_image(image, self.signature_inputs["Image"]["shape"]) | |
| feed = {self.signature_inputs["Image"]["name"]: [img]} | |
| output_names = [self.signature_outputs[key]["name"] for key in self.signature_outputs] | |
| outputs = self.session.run(output_names=output_names, input_feed=feed) | |
| return self.process_output(outputs) | |
| def process_image(self, image: Image.Image, input_shape: list) -> np.ndarray: | |
| """Resize and normalize the image.""" | |
| width, height = image.size | |
| if image.mode != "RGB": | |
| image = image.convert("RGB") | |
| 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 | |
| image = image.crop((left, top, right, bottom)) | |
| input_width, input_height = input_shape[1:3] | |
| image = image.resize((input_width, input_height)) | |
| image = np.asarray(image) / 255.0 | |
| return image.astype(np.float32) | |
| def process_output(self, outputs: list) -> dict: | |
| """Format the model output.""" | |
| out_keys = ["label", "confidence"] | |
| results = {key: outputs[i].tolist()[0] for i, key in enumerate(self.signature_outputs)} | |
| confs = results["Confidences"] | |
| labels = self.signature["classes"]["Label"] | |
| output = [dict(zip(out_keys, group)) for group in zip(labels, confs)] | |
| return {"predictions": sorted(output, key=lambda x: x["confidence"], reverse=True)} | |
| EXPORT_MODEL_VERSION = 1 | |
| model = ONNXModel(dir_path="model.onnx") | |
| model.load() | |
| def predict(image): | |
| """Run inference on the given 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.Image(image_mode="RGB") | |
| outputs = gr.JSON() | |
| description = ( | |
| "This is a web interface for the Naked Detector model. " | |
| "Upload an image and get predictions for the presence of nudity.\n" | |
| "Model and website created by KUO SUKO, C110156115 NKUST." | |
| ) | |
| interface = gr.Interface( | |
| fn=predict, | |
| inputs=inputs, | |
| outputs=outputs, | |
| title="Naked Detector", | |
| description=description | |
| ) | |
| interface.launch() | |
| # this is changed by ChatGPT, if it run like a shit. don't blame me >< |