Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import torch | |
| from PIL import Image | |
| from transformers import AutoModelForImageClassification, ViTImageProcessor | |
| # Load model | |
| model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection") | |
| processor = ViTImageProcessor.from_pretrained('Falconsai/nsfw_image_detection') | |
| # Function to make predictions | |
| def classify_image(img): | |
| with torch.no_grad(): | |
| inputs = processor(images=img, return_tensors="pt") | |
| outputs = model(**inputs) | |
| logits = outputs.logits | |
| predicted_label = logits.argmax(-1).item() | |
| return model.config.id2label[predicted_label] | |
| # Gradio Interface | |
| iface = gr.Interface( | |
| fn=classify_image, | |
| inputs="image", | |
| outputs="text", | |
| live=True # Remove the 'interpretation' argument | |
| ) | |
| # Launch the Gradio Interface | |
| iface.launch() |