Commit a592c144 authored by Anjith GEORGE's avatar Anjith GEORGE Committed by Anjith GEORGE

Added a MultiSpectral extension for VGG16 for PAD

parent 5335c905
import torch
from torch import nn
from torchvision import models
class DeepMSPAD(nn.Module):
""" Deep multispectral PAD algorithm
Attributes:
pretrained: bool
if set `True` loads the pretrained vgg16 model.
vgg: :py:class:`torch.nn.Module`
The VGG16 model
relu: :py:class:`torch.nn.Module`
ReLU activation
enc: :py:class:`torch.nn.Module`
Uses the layers for feature extraction
linear1: :py:class:`torch.nn.Module`
Fully connected layer
linear2: :py:class:`torch.nn.Module`
Fully connected layer
dropout: :py:class:`torch.nn.Module`
Dropout layer
sigmoid: :py:class:`torch.nn.Module`
Sigmoid activation
"""
def __init__(self, pretrained=True, num_channels=4):
""" Init method
Parameters
----------
pretrained: bool
if set `True` loads the pretrained vgg16 model.
num_channels: int
Number of channels in the input
"""
super(DeepMSPAD, self).__init__()
vgg = models.vgg16(pretrained=pretrained)
features = list(vgg.features.children())
features[0]=nn.Conv2d(num_channels, 64, kernel_size=(3, 3), stride=(1, 1), padding =(1, 1))
self.enc = nn.Sequential(*features)
self.linear1=nn.Linear(25088,256)
self.relu=nn.ReLU()
self.dropout= nn.Dropout(p=0.5)
self.linear2=nn.Linear(256,1)
self.sigmoid= nn.Sigmoid()
def forward(self, x):
""" Propagate data through the network
Parameters
----------
x: :py:class:`torch.Tensor`
The data to forward through the network
Returns
-------
x: :py:class:`torch.Tensor`
The last layer of the network
"""
enc = self.enc(x)
x=enc.view(-1,25088)
x=self.linear1(x)
x=self.relu(x)
x=self.dropout(x)
x=self.linear2(x)
x=self.sigmoid(x)
return x
......@@ -6,6 +6,7 @@ from .LightCNN import LightCNN29v2
from .MCCNN import MCCNN
from .MCCNNv2 import MCCNNv2
from .FASNet import FASNet
from .DeepMSPAD import DeepMSPAD
from .DCGAN import DCGAN_generator
from .DCGAN import DCGAN_discriminator
......
......@@ -84,6 +84,14 @@ def test_architectures():
output = net.forward(t)
assert output.shape == torch.Size([1, 1])
#DeepMSPAD
a = numpy.random.rand(1, 8, 224, 224).astype("float32")
t = torch.from_numpy(a)
from ..architectures import DeepMSPAD
net = FASNet(pretrained=False, num_channels=8)
output = net.forward(t)
assert output.shape == torch.Size([1, 1])
# DCGAN
d = numpy.random.rand(1, 3, 64, 64).astype("float32")
t = torch.from_numpy(d)
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment