import gradio as gr
from torchvision.models import resnet50
import torch
import torch.nn as nn
from torch.nn import functional as F
from torchvision import transforms
import numpy as np
import os
import random
def seed_torch(seed=1029):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.enabled = False
seed_torch(100)

NPRmodel = resnet50()
NPRmodel.fc1 = nn.Linear(512, 1)
NPRmodel.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
del NPRmodel.layer3, NPRmodel.layer4, NPRmodel.fc 

from collections import OrderedDict
from copy import deepcopy

NPR_url = 'https://raw.githubusercontent.com/chuangchuangtan/NPR-DeepfakeDetection/main/model_epoch_last_3090.pth'
NPRmodel.load_state_dict(torch.hub.load_state_dict_from_url(NPR_url, map_location='cpu'), strict=True)
trans = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])

def interpolate(img, factor):
    return F.interpolate(F.interpolate(img, scale_factor=factor, mode='nearest', recompute_scale_factor=True), scale_factor=1/factor, mode='nearest', recompute_scale_factor=True)

def predict_image(img):
    img = trans( img.convert('RGB') ).unsqueeze(0)
    _, c, w, h = img.shape
    if w%2 == 1: img = img[:, :, :-1,:  ]
    if h%2 == 1: img = img[:, :, :  ,:-1]
    NPR = img - interpolate(img, 0.5)
    with torch.no_grad():
        x   = NPRmodel.conv1(NPR*2.0/3.0)
        x   = NPRmodel.bn1(x)
        x   = NPRmodel.relu(x)
        x   = NPRmodel.maxpool(x)
        x   = NPRmodel.layer1(x)
        x   = NPRmodel.layer2(x).mean(dim=(2,3), keepdim=False)
        x   = NPRmodel.fc1(x)
        pred = x.sigmoid().cpu().numpy()
    return {'Fake Image': float(pred), 'Real Image': 1.0 - float(pred)}


# Define the Gradio Interface with the desired title and description
title = "<p>Generalizable Deepfake Detection Neighboring Pixel Relationships(NPR)-CVPR2024<p>"
description_html = """
<p> <strong>Rethinking the Up-Sampling Operations in CNN-based Generative Network for Generalizable Deepfake Detection * CVPR2024</strong><p>
<p> This model was only trained on ProGAN for detecting deepfake images. </p>
<p> The model achieved an accuracy of <strong>92.2%</strong> on the 28 generation models.</p>
<p> Upload a image or pick from the samples below to test model accuracy</p>
<p><a href='https://github.com/chuangchuangtan/NPR-DeepfakeDetection' target='_blank'>NPR Github Code</a></p>
"""

# Define example images and their true labels for users to choose from
example_data = ['midjourney1.png', 'CelebAHQ_00000110.png']
custom_css = """
div {background-color: whitesmoke;}
"""

gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=2),
    title=title,
    description=description_html,
    allow_flagging='never',
    examples=example_data,
    css=custom_css
).launch()
# ).launch(server_name="0.0.0.0", server_port=80)