Create app.py

app.py

@@ -0,0 +1,32 @@
+import gradio as gr
+import torch
+import timm
+from PIL import Image
+
+model = timm.create_model("hf_hub:Marqo/nsfw-image-detection-384", pretrained=True).eval()
+data_config = timm.data.resolve_model_data_config(model)
+transforms = timm.data.create_transform(**data_config, is_training=False)
+class_names = model.pretrained_cfg["label_names"]
+
+@torch.inference_mode()
+def predict(image: Image.Image):
+    tensor = transforms(image).unsqueeze(0)
+    probs  = model(tensor).softmax(dim=-1).cpu().flatten()
+    top_id = int(probs.argmax())
+    top_label = class_names[top_id]
+    probs_dict = {class_names[i]: float(p) for i, p in enumerate(probs)}
+    return top_label, probs_dict
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=[
+        gr.Label(label="Top prediction"),
+        gr.Label(label="All probabilities", num_top_classes=len(class_names)),
+    ],
+    title="NSFW Image Detection",
+    description="Drag & drop an image to see the predicted class",
+)
+
+if __name__ == "__main__":
+    demo.launch(show_error=True)