iris-backend / pad_model.py
trretretret's picture
Fix static structure for HF deployment
9232fb9
Raw
History Blame Contribute Delete
3.51 kB
# import torch
# import torch.nn as nn
# import torch.nn.functional as F
# from PIL import Image
# import numpy as np
# import cv2
# from torchvision import transforms
# from config import PAD_MODEL_PATH
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# class DNetPAD(nn.Module):
# def __init__(self, num_classes=2):
# super().__init__()
# self.conv1=nn.Conv2d(3,32,3,1); self.bn1=nn.BatchNorm2d(32)
# self.conv2=nn.Conv2d(32,64,3,1); self.bn2=nn.BatchNorm2d(64)
# self.conv3=nn.Conv2d(64,128,3,1); self.bn3=nn.BatchNorm2d(128)
# self.conv4=nn.Conv2d(128,256,3,1); self.bn4=nn.BatchNorm2d(256)
# self.pool = nn.MaxPool2d(2,2); self.dropout=nn.Dropout(0.5)
# self.gap = nn.AdaptiveAvgPool2d(1)
# self.fc1 = nn.Linear(256,128); self.fc2=nn.Linear(128,64); self.fc3=nn.Linear(64,num_classes)
# self.last_conv_output = None
# def forward(self, x):
# x=F.relu(self.bn1(self.conv1(x))); x=self.pool(x)
# x=F.relu(self.bn2(self.conv2(x))); x=self.pool(x)
# x=F.relu(self.bn3(self.conv3(x))); x=self.pool(x)
# x=F.relu(self.bn4(self.conv4(x))); self.last_conv_output=x; x=self.pool(x)
# x=self.gap(x); x=x.view(x.size(0),-1)
# x=F.relu(self.fc1(x)); x=self.dropout(x)
# x=F.relu(self.fc2(x)); x=self.dropout(x)
# x=self.fc3(x)
# return x
# # pad_model = DNetPAD().to(device)
# # pad_model.load_state_dict(torch.load(PAD_MODEL_PATH,map_location=device))
# # pad_model.eval()
# pad_model = DNetPAD().to(device)
# try:
# pad_model.load_state_dict(torch.load(PAD_MODEL_PATH,map_location=device))
# pad_model.eval()
# except Exception as e:
# print("Warning: PAD model could not be loaded:", e)
# pad_transform = transforms.Compose([
# transforms.Resize((224,224)),
# transforms.ToTensor(),
# transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
# ])
# def pad_predict_and_explain(image_path):
# img = Image.open(image_path).convert("RGB")
# input_tensor = pad_transform(img).unsqueeze(0).to(device)
# with torch.no_grad():
# outputs = pad_model(input_tensor)
# probs = torch.softmax(outputs, dim=1)
# pred_class = torch.argmax(probs, dim=1).item()
# confidence = probs[0, pred_class].item()
# label = "Fake" if pred_class==1 else "Real"
# # Grad-CAM overlay (simplified)
# try:
# pad_model.zero_grad()
# outputs[0,pred_class].backward(retain_graph=True)
# activations = pad_model.last_conv_output
# gradients = activations.grad if activations.grad is not None else torch.ones_like(activations)
# pooled_gradients = torch.mean(gradients,dim=[0,2,3])
# for i in range(activations.shape[1]):
# activations[:,i,:,:] *= pooled_gradients[i]
# heatmap = torch.mean(activations, dim=1).squeeze().cpu().numpy()
# heatmap = np.maximum(heatmap,0)
# if np.max(heatmap)>0: heatmap/=np.max(heatmap)
# heatmap=cv2.resize(heatmap,(224,224))
# heatmap_colored=cv2.applyColorMap(np.uint8(255*heatmap),cv2.COLORMAP_JET)
# heatmap_colored=cv2.cvtColor(heatmap_colored,cv2.COLOR_BGR2RGB)
# overlay = cv2.addWeighted(np.array(img.resize((224,224))),0.6,heatmap_colored,0.4,0)
# except:
# overlay = np.array(img.resize((224,224)))
# return {"prediction": label, "confidence": round(float(confidence),3), "heatmap_image": overlay}