def __init__(self, lab_version, input_size=128):
super(ColorCNN_v2, self).__init__()
MIDLEVEL_FEATURE_SIZE = 128
if lab_version == 1:
self.final = nn.Tanh()
elif lab_version == 2:
self.final = nn.ReLU()
## First half: ResNet
resnet = models.resnet18(num_classes=365)
# Change first conv layer to accept single-channel (grayscale) input
resnet.conv1.weight = nn.Parameter(
resnet.conv1.weight.sum(dim=1).unsqueeze(1))
# Extract midlevel features from ResNet-gray
self.midlevel_resnet = nn.Sequential(*list(resnet.children())[0:6])
self.upsampling = Upsample(scale_factor=2, mode='nearest')
## Second half: Upsampling
self.upsample = nn.Sequential(
nn.Conv2d(MIDLEVEL_FEATURE_SIZE,
128,
kernel_size=3,
stride=1,
padding=1), nn.BatchNorm2d(128), nn.ReLU(),
self.upsampling,
nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(64), nn.ReLU(),
nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(64), nn.ReLU(), self.upsampling,
nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(32), self.final,
nn.Conv2d(32, 2, kernel_size=3, stride=1,
padding=1), self.upsampling)
def __init__(self, lab_version):
super(ColorCNN_v1, self).__init__()
self.relu = nn.ReLU()
if lab_version == 1:
self.final = nn.Tanh()
elif lab_version == 2:
self.final = nn.Sigmoid()
self.upsampling = Upsample(scale_factor=2, mode='nearest')
self.conv1 = nn.Conv2d(in_channels=1, out_channels=64, kernel_size=(3,3), stride=1, padding=1, bias=True)
self.conv2 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=(3,3), stride=2, padding=1, bias=True)
self.conv3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=(3,3), stride=1, padding=1, bias=True)
self.conv4 = nn.Conv2d(in_channels=128, out_channels=128, kernel_size=(3,3), stride=2, padding=1, bias=True)
self.conv5 = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=(3,3), stride=1, padding=1, bias=True)
self.conv6 = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=(3,3), stride=2, padding=1, bias=True)
self.conv7 = nn.Conv2d(in_channels=256, out_channels=512, kernel_size=(3,3), stride=1, padding=1, bias=True)
self.conv8 = nn.Conv2d(in_channels=512, out_channels=256, kernel_size=(3,3), stride=1, padding=1, bias=True)
self.conv9 = nn.Conv2d(in_channels=256, out_channels=128, kernel_size=(3,3), stride=1, padding=1, bias=True)
self.conv10 = nn.Conv2d(in_channels=128, out_channels=64, kernel_size=(3,3), stride=1, padding=1, bias=True)
self.conv11 = nn.Conv2d(in_channels=64, out_channels=32, kernel_size=(3,3), stride=1, padding=1, bias=True)
self.conv12 = nn.Conv2d(in_channels=32, out_channels=2, kernel_size=(3,3), stride=1, padding=1, bias=True)
self.color = nn.Sequential(
self.conv1,
self.relu,
self.conv2,
self.relu,
self.conv3,
self.relu,
self.conv4,
self.relu,
self.conv5,
self.relu,
self.conv6,
self.relu,
self.conv7,
self.relu,
self.conv8,
self.relu,
self.conv9,
self.relu,
self.upsampling,
self.conv10,
self.relu,
self.upsampling,
self.conv11,
self.relu,
self.conv12,
self.final,
self.upsampling
)