Spaces:
Build error
Build error
| import gradio as gr | |
| from transformers import ViTFeatureExtractor, ViTForImageClassification | |
| import torch | |
| import requests | |
| # Load the pre-trained ViT model and feature extractor | |
| model = ViTForImageClassification.from_pretrained("umitkantar/ViTBasePatch16Art") | |
| feature_extractor = ViTFeatureExtractor.from_pretrained("umitkantar/ViTBasePatch16Art") | |
| # Define the prediction function | |
| def ai_image_detector(image): | |
| # Preprocess the input image | |
| inputs = feature_extractor(images=image, return_tensors="pt") | |
| # Make a prediction using the loaded model | |
| outputs = model(**inputs) | |
| predicted_class_idx = torch.argmax(outputs.logits).item() | |
| # Calculate the confidence score (percentage) | |
| confidence_score = torch.softmax(outputs.logits, dim=1)[0][predicted_class_idx].item() * 100 | |
| return {f"AI-generated Probability (%): {confidence_score:.2f}": confidence_score} | |
| # Set up the Gradio interface | |
| image_input = gr.inputs.Image() | |
| gr.Interface(ai_image_detector, image_input, "label", capture_session=True).launch() | |