File size: 1,635 Bytes
2d2db50 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | 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)
|