aruntandra's picture
Upload 4 files
edba0cc verified
Raw
History Blame Contribute Delete
1.87 kB
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
def ensure_gray(image: Image.Image) -> Image.Image:
if image.mode != "L":
image = image.convert("L")
return image
trnscm = transforms.Compose([
transforms.Lambda(ensure_gray),
transforms.Resize((100, 100)),
transforms.ToTensor(),
])
class Siamese(nn.Module):
def __init__(self):
super().__init__()
self.cnn1 = nn.Sequential(
nn.ReflectionPad2d(1),
nn.Conv2d(1, 4, kernel_size=3),
nn.ReLU(inplace=True),
nn.BatchNorm2d(4),
nn.ReflectionPad2d(1),
nn.Conv2d(4, 8, kernel_size=3),
nn.ReLU(inplace=True),
nn.BatchNorm2d(8),
nn.ReflectionPad2d(1),
nn.Conv2d(8, 8, kernel_size=3),
nn.ReLU(inplace=True),
nn.BatchNorm2d(8),
)
self.fc1 = nn.Sequential(
nn.Linear(8 * 100 * 100, 500),
nn.ReLU(inplace=True),
nn.Linear(500, 500),
nn.ReLU(inplace=True),
nn.Linear(500, 5),
)
def forward_once(self, x):
out = self.cnn1(x)
out = out.view(out.size(0), -1)
out = self.fc1(out)
return out
def forward(self, x1, x2):
return self.forward_once(x1), self.forward_once(x2)
class FaceClassifier(nn.Module):
def __init__(self, input_dim=5, num_classes=7):
super().__init__()
self.fc = nn.Sequential(
nn.Linear(input_dim, 128),
nn.ReLU(inplace=True),
nn.Dropout(0.5),
nn.Linear(128, num_classes),
)
def forward(self, x):
return self.fc(x)
# Update to match your captured_face_images ImageFolder order
classes = ["person1","person2","person3","person4","person5","person6","person7"]