Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import torch | |
| import torch.nn as nn | |
| import torchvision.transforms as transforms | |
| from PIL import Image | |
| import torch.nn.init as init | |
| class Fire(nn.Module): | |
| def __init__(self, inplanes: int, squeeze_planes: int, expand1x1_planes: int, expand3x3_planes: int) -> None: | |
| super().__init__() | |
| self.inplanes = inplanes | |
| self.squeeze = nn.Conv2d(inplanes, squeeze_planes, kernel_size=1) | |
| self.squeeze_activation = nn.ReLU(inplace=True) | |
| self.expand1x1 = nn.Conv2d(squeeze_planes, expand1x1_planes, kernel_size=1) | |
| self.expand1x1_activation = nn.ReLU(inplace=True) | |
| self.expand3x3 = nn.Conv2d(squeeze_planes, expand3x3_planes, kernel_size=3, padding=1) | |
| self.expand3x3_activation = nn.ReLU(inplace=True) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| x = self.squeeze_activation(self.squeeze(x)) | |
| return torch.cat( | |
| [self.expand1x1_activation(self.expand1x1(x)), self.expand3x3_activation(self.expand3x3(x))], 1 | |
| ) | |
| class SqueezeNet(nn.Module): | |
| def __init__(self, version: str = "1_0", num_classes: int = 1000, dropout: float = 0.5) -> None: | |
| super().__init__() | |
| # _log_api_usage_once(self) | |
| self.num_classes = num_classes | |
| if version == "1_0": | |
| self.features = nn.Sequential( | |
| nn.Conv2d(3, 96, kernel_size=7, stride=2), | |
| nn.ReLU(inplace=True), | |
| nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), | |
| Fire(96, 16, 64, 64), | |
| Fire(128, 16, 64, 64), | |
| Fire(128, 32, 128, 128), | |
| nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), | |
| Fire(256, 32, 128, 128), | |
| Fire(256, 48, 192, 192), | |
| Fire(384, 48, 192, 192), | |
| Fire(384, 64, 256, 256), | |
| nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), | |
| Fire(512, 64, 256, 256), | |
| ) | |
| elif version == "middle": # 0.78 mb | |
| self.features = nn.Sequential( | |
| nn.Conv2d(3, 32, kernel_size=3, stride=2), | |
| nn.ReLU(inplace=True), | |
| nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), | |
| Fire(32, 8, 32, 32), | |
| Fire(64, 8, 32, 32), | |
| nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), | |
| Fire(64, 16, 64, 64), | |
| Fire(128, 16, 64, 64), | |
| nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), | |
| Fire(128, 24, 96, 96), | |
| Fire(192, 24, 96, 96), | |
| Fire(192, 32, 128, 128), | |
| Fire(256, 32, 128, 128), | |
| ) | |
| else: | |
| # FIXME: Is this needed? SqueezeNet should only be called from the | |
| # FIXME: squeezenet1_x() functions | |
| # FIXME: This checking is not done for the other models | |
| raise ValueError(f"Unsupported SqueezeNet version {version}: 1_0 or 1_1 expected") | |
| # Final convolution is initialized differently from the rest | |
| # 512 | |
| final_conv = nn.Conv2d(256, self.num_classes, kernel_size=1) | |
| self.classifier = nn.Sequential( | |
| nn.Dropout(p=dropout), final_conv, nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d((1, 1)) | |
| ) | |
| for m in self.modules(): | |
| if isinstance(m, nn.Conv2d): | |
| if m is final_conv: | |
| init.normal_(m.weight, mean=0.0, std=0.01) | |
| else: | |
| init.kaiming_uniform_(m.weight) | |
| if m.bias is not None: | |
| init.constant_(m.bias, 0) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| x = self.features(x) | |
| x = self.classifier(x) | |
| return torch.flatten(x, 1) | |
| class SN(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.model = SqueezeNet(version="middle", num_classes=2) | |
| def forward(self, x): | |
| return self.model(x) | |
| def predict(image): | |
| model = SN() | |
| model.load_state_dict(torch.load("./liveness_1M_model_0.8740054619288524 .ckpt", map_location=torch.device('cpu'))) | |
| im = Image.open(image) | |
| transform1 = transforms.Compose([transforms.Resize((512, 512)), | |
| transforms.ToTensor()]) | |
| img = transform1(im).unsqueeze(0) | |
| my_softmax = nn.Softmax(dim=1) | |
| with torch.no_grad(): | |
| y_hat = model(img) | |
| liveness_score = float(my_softmax(y_hat)[0][1]) | |
| res = {"fake": liveness_score, "real": 1 - liveness_score} | |
| return res | |
| gr.Interface( | |
| predict, | |
| inputs=gr.inputs.Image(label="Upload an image", type="filepath"), | |
| outputs=gr.outputs.Label(num_top_classes=2), | |
| title="Real or Fake", examples=["./2022-10-12 16.52.56.jpg", "./2022-10-12 16.54.52.jpg", "./1724477482.jpeg"] | |
| ).launch() |