David310/Detect_AI-generated_Image

1import gradio as gr2import os3import csv4from models import get_model5import torch6import torchvision.transforms as transforms7import torch.utils.data8import numpy as np9import sys10from PIL import Image11# from detect_one_image import detect_one_image12 13MEAN = {14    "imagenet":[0.485, 0.456, 0.406],15    "clip":[0.48145466, 0.4578275, 0.40821073]16}17 18STD = {19    "imagenet":[0.229, 0.224, 0.225],20    "clip":[0.26862954, 0.26130258, 0.27577711]21}22 23 24def detect_one_image(model, image):25 26    """27    model = get_model('CLIP:ViT-L/14')28    state_dict = torch.load(ckpt, map_location='cpu')29    model.fc.load_state_dict(state_dict)30    print ("Model loaded..")31    model.eval()32    model.cuda()33    """34    # img = Image.open(image_path).convert("RGB")35    """36    if jpeg_quality is not None:37        img = png2jpg(img, jpeg_quality)38    """39    transform = transforms.Compose([40            transforms.ToTensor(),41            transforms.CenterCrop(224),42            transforms.Normalize( mean=MEAN['clip'], std=STD['clip'] ),43        ])44    img = transform(image)45    img = img.to('cuda:0')46 47    detection_output = model(img)48    output = torch.sigmoid(detection_output)49 50    return output51 52def detect(image):53    # print(type(image))54    model = get_model('CLIP:ViT-L/14')55    state_dict = torch.load('./pretrained_weights/fc_weights.pth', map_location='cpu')56    model.fc.load_state_dict(state_dict)57    # model.load_state_dict(state_dict)58    # print ("Model loaded..")59    model.eval()60    model.cuda()61    output_tensor = detect_one_image(model, image)62    ai_likelihood = (100*output_tensor).item()63    return "The image is " + str(ai_likelihood) + r" % likely to be AI-generated."64 65demo = gr.Interface(66    fn=detect,67    inputs=["image"],68    outputs=["text"],69)70 71demo.launch()72