Views
No views yet
1# Custom Head Implementation
2model.fc = nn.Sequential(
3 nn.Dropout(0.1),
4 nn.Linear(512, 2) # Mapping features to ['Forged', 'Original']
5)model.pt artifact is pushed to the cloud for production readiness.1import torch
2from torchvision import transforms
3from PIL import Image
4
5# 1. Setup Device
6device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
7
8# 2. Load the entire model object
9# Ensure the 'model.pt' file is in your working directory
10model = torch.load("model.pt", map_location=device)
11model.eval()
12
13# 3. Define Prediction Function
14def predict(image_path):
15 img = Image.open(image_path).convert('RGB')
16
17 # Matches the transformation pipeline used during training
18 transform = transforms.Compose([
19 transforms.Resize((224, 224)),
20 transforms.ToTensor(),
21 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
22 ])
23
24 img_tensor = transform(img).unsqueeze(0).to(device)
25
26 with torch.no_grad():
27 output = model(img_tensor)
28 probs = torch.softmax(output, dim=1)
29 prediction = torch.argmax(probs, dim=1).item()
30
31 labels = ['Forged', 'Original']
32 return labels[prediction]
33
34# Example run:
35# print(f"Verdict: {predict('test_signature.png')}")