from fastai.vision.all import load_learner

import gradio as gr

from fastai.vision.all import *
import numpy as np
import torch

dummy_img = PILImage.create(np.zeros((224,224,3), dtype=np.uint8))

dblock = DataBlock(
    blocks=(ImageBlock, CategoryBlock(vocab=['class_0', 'class_1'])),
    get_items=lambda _: [dummy_img],
    get_y=lambda _: 'class_0',
    item_tfms=Resize(224)
)

dls = dblock.dataloaders(source=None, bs=1, device='cpu')

learn = vision_learner(
    dls,
    resnet18,
    n_out=2,
    pretrained=False,
    normalize=False
)

learn.model.load_state_dict(
    torch.load("toai.bin", map_location="cpu")
)

learn.model.eval()

from fastai.vision.all import PILImage
from PIL import Image
import requests
from io import BytesIO

categories = ['AI Generated', 'Not AI Generated']

def classify_image(img):

    fastai_img = PILImage.create(img)
    pred, pred_idx, probs = learn.predict(fastai_img)

    return dict(zip(categories, map(float, probs)))

interface = gr.Interface(fn=classify_image, inputs=gr.Image(type="pil", image_mode="RGB"), 
    outputs=gr.Label(),
    examples=["https://huggingface.co/datasets/brawl7787/example_img/resolve/main/william_D.jpeg", 
              "https://huggingface.co/datasets/brawl7787/example_img/resolve/main/william_D_dall-e.jpeg",
              "https://huggingface.co/datasets/brawl7787/example_img/resolve/main/william_D_gemini.jpg"],
    cache_examples=False,
    title="AI or Not AI Image Classifier",
    description="Upload an image to determine if it is AI generated or not.")
interface.launch(inline=True)