|
|
import gradio as gr |
|
|
import numpy as np |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.optim as optim |
|
|
import torch.nn.functional as F |
|
|
import cv2 |
|
|
import PIL.Image |
|
|
from scipy.interpolate import griddata |
|
|
import matplotlib.pyplot as plt |
|
|
from utils import azi_diff |
|
|
|
|
|
class AttentionBlock(nn.Module): |
|
|
def __init__(self, input_dim, num_heads, ff_dim, rate=0.2): |
|
|
super(AttentionBlock, self).__init__() |
|
|
self.attention = nn.MultiheadAttention(embed_dim=input_dim, num_heads=num_heads) |
|
|
self.dropout1 = nn.Dropout(rate) |
|
|
self.layer_norm1 = nn.LayerNorm(input_dim) |
|
|
|
|
|
self.ffn = nn.Sequential( |
|
|
nn.Linear(input_dim, ff_dim), |
|
|
nn.ReLU(), |
|
|
nn.Dropout(rate), |
|
|
nn.Linear(ff_dim, input_dim), |
|
|
nn.Dropout(rate) |
|
|
) |
|
|
self.layer_norm2 = nn.LayerNorm(input_dim) |
|
|
|
|
|
def forward(self, x): |
|
|
attn_output, _ = self.attention(x, x, x) |
|
|
attn_output = self.dropout1(attn_output) |
|
|
out1 = self.layer_norm1(attn_output + x) |
|
|
|
|
|
ffn_output = self.ffn(out1) |
|
|
out2 = self.layer_norm2(ffn_output + out1) |
|
|
return out2 |
|
|
|
|
|
class TextureContrastClassifier(nn.Module): |
|
|
def __init__(self, input_shape, num_heads=4, key_dim=64, ff_dim=256, rate=0.1): |
|
|
super(TextureContrastClassifier, self).__init__() |
|
|
input_dim = input_shape[1] |
|
|
self.rich_attention_block = AttentionBlock(input_dim, num_heads, ff_dim, rate) |
|
|
self.rich_dense = nn.Sequential( |
|
|
nn.Linear(input_dim, 128), |
|
|
nn.ReLU(), |
|
|
nn.Dropout(0.5) |
|
|
) |
|
|
self.poor_attention_block = AttentionBlock(input_dim, num_heads, ff_dim, rate) |
|
|
self.poor_dense = nn.Sequential( |
|
|
nn.Linear(input_dim, 128), |
|
|
nn.ReLU(), |
|
|
nn.Dropout(0.5) |
|
|
) |
|
|
self.fc = nn.Sequential( |
|
|
nn.Linear(128 * input_shape[0], 256), |
|
|
nn.ReLU(), |
|
|
nn.Dropout(0.5), |
|
|
nn.Linear(256, 128), |
|
|
nn.ReLU(), |
|
|
nn.Dropout(0.5), |
|
|
nn.Linear(128, 64), |
|
|
nn.ReLU(), |
|
|
nn.Dropout(0.5), |
|
|
nn.Linear(64, 32), |
|
|
nn.ReLU(), |
|
|
nn.Dropout(0.5), |
|
|
nn.Linear(32, 16), |
|
|
nn.ReLU(), |
|
|
nn.Dropout(0.5), |
|
|
nn.Linear(16, 1), |
|
|
nn.Sigmoid() |
|
|
) |
|
|
|
|
|
def forward(self, rich_texture, poor_texture): |
|
|
rich_texture = rich_texture.permute(1, 0, 2) |
|
|
poor_texture = poor_texture.permute(1, 0, 2) |
|
|
rich_attention = self.rich_attention_block(rich_texture) |
|
|
rich_attention = rich_attention.permute(1, 0, 2) |
|
|
rich_features = self.rich_dense(rich_attention) |
|
|
poor_attention = self.poor_attention_block(poor_texture) |
|
|
poor_attention = poor_attention.permute(1, 0, 2) |
|
|
poor_features = self.poor_dense(poor_attention) |
|
|
difference = rich_features - poor_features |
|
|
difference = difference.view(difference.size(0), -1) |
|
|
output = self.fc(difference) |
|
|
return output |
|
|
|
|
|
input_shape = (128, 256) |
|
|
model = TextureContrastClassifier(input_shape) |
|
|
model.load_state_dict(torch.load('./model_epoch_36.pth', map_location=torch.device('cpu'))) |
|
|
|
|
|
def inference(image, model): |
|
|
predictions = [] |
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
model.to(device) |
|
|
model.eval() |
|
|
tmp = azi_diff(image, patch_num=128, N=256) |
|
|
rich = tmp["total_emb"][0] |
|
|
poor = tmp["total_emb"][1] |
|
|
rich_texture_tensor = torch.tensor(rich, dtype=torch.float32).unsqueeze(0).to(device) |
|
|
poor_texture_tensor = torch.tensor(poor, dtype=torch.float32).unsqueeze(0).to(device) |
|
|
with torch.no_grad(): |
|
|
output = model(rich_texture_tensor, poor_texture_tensor) |
|
|
prediction = output.cpu().numpy().flatten()[0] |
|
|
return prediction |
|
|
|
|
|
|
|
|
def predict(image): |
|
|
prediction = inference(image, model) |
|
|
return f"{prediction * 100:.2f}% chance AI-generated" |
|
|
|
|
|
gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs="text").launch() |
|
|
|
