import gradio as gr
from transformers import AutoFeatureExtractor, AutoModelForImageClassification, CLIPVisionModel
import torch
from PIL import Image
import json
import numpy as np

# Load pre-trained model and feature extractor
model_name = "mrsarthakgupta/openclipvisionmodel"
feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
model = CLIPVisionModel.from_pretrained(model_name)

matrix = torch.Tensor(np.array(json.load(open('linear_params.json', 'r'))["0"]))

def classify_image(image):
    # Preprocess the image
    inputs = feature_extractor(images=image, return_tensors="pt")
    
    # Make prediction
    with torch.no_grad():
        outputs = model(**inputs)
    
    results = torch.matmul(outputs.pooler_output.squeeze(), matrix.squeeze())
    return {'Our confidence of this image being AI generated': torch.nn.functional.sigmoid(results + torch.Tensor([0.9791784]))}


# Create Gradio interface
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=1),
    title="Detect AI generated image",
    description="Upload an image to find the probability of it being AI generated. While we're certainly not perfect, but we're working really hard to get a more accurate classifier!"
)

# Launch the app
iface.launch()