| 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) | |