cleanup

face_swapper/src/networks/encoder.py

@@ -0,0 +1,67 @@
+import torch
+from torch import Tensor, nn as nn
+
+from face_swapper.src.types import TargetAttributes, VisionTensor
+
+
+class Upsample(nn.Module):
+	def __init__(self, input_channels : int, output_channels : int) -> None:
+		super(Upsample, self).__init__()
+		self.deconv = nn.ConvTranspose2d(in_channels = input_channels, out_channels = output_channels, kernel_size = 4, stride = 2, padding = 1, bias = False)
+		self.batch_norm = nn.BatchNorm2d(output_channels)
+		self.leaky_relu = nn.LeakyReLU(0.1, inplace = True)
+
+	def forward(self, temp : Tensor, skip_tensor : Tensor) -> Tensor:
+		temp = self.deconv(temp)
+		temp = self.batch_norm(temp)
+		temp = self.leaky_relu(temp)
+		return torch.cat((temp, skip_tensor), dim = 1)
+
+
+class DownSample(nn.Module):
+	def __init__(self, input_channels : int, output_channels : int) -> None:
+		super(DownSample, self).__init__()
+		self.conv = nn.Conv2d(in_channels = input_channels, out_channels = output_channels, kernel_size = 4, stride = 2, padding = 1, bias = False)
+		self.batch_norm = nn.BatchNorm2d(output_channels)
+		self.leaky_relu = nn.LeakyReLU(0.1, inplace = True)
+
+	def forward(self, temp : Tensor) -> Tensor:
+		temp = self.conv(temp)
+		temp = self.batch_norm(temp)
+		temp = self.leaky_relu(temp)
+		return temp
+
+
+class UNet(nn.Module):
+	def __init__(self) -> None:
+		super(UNet, self).__init__()
+		self.downsampler_1 = DownSample(3, 32)
+		self.downsampler_2 = DownSample(32, 64)
+		self.downsampler_3 = DownSample(64, 128)
+		self.downsampler_4 = DownSample(128, 256)
+		self.downsampler_5 = DownSample(256, 512)
+		self.downsampler_6 = DownSample(512, 1024)
+		self.bottleneck = DownSample(1024, 1024)
+		self.upsampler_1 = Upsample(1024, 1024)
+		self.upsampler_2 = Upsample(2048, 512)
+		self.upsampler_3 = Upsample(1024, 256)
+		self.upsampler_4 = Upsample(512, 128)
+		self.upsampler_5 = Upsample(256, 64)
+		self.upsampler_6 = Upsample(128, 32)
+
+	def forward(self, target : VisionTensor) -> TargetAttributes:
+		downsample_feature_1 = self.downsampler_1(target)
+		downsample_feature_2 = self.downsampler_2(downsample_feature_1)
+		downsample_feature_3 = self.downsampler_3(downsample_feature_2)
+		downsample_feature_4 = self.downsampler_4(downsample_feature_3)
+		downsample_feature_5 = self.downsampler_5(downsample_feature_4)
+		downsample_feature_6 = self.downsampler_6(downsample_feature_5)
+		bottleneck_output = self.bottleneck(downsample_feature_6)
+		upsample_feature_1 = self.upsampler_1(bottleneck_output, downsample_feature_6)
+		upsample_feature_2 = self.upsampler_2(upsample_feature_1, downsample_feature_5)
+		upsample_feature_3 = self.upsampler_3(upsample_feature_2, downsample_feature_4)
+		upsample_feature_4 = self.upsampler_4(upsample_feature_3, downsample_feature_3)
+		upsample_feature_5 = self.upsampler_5(upsample_feature_4, downsample_feature_2)
+		upsample_feature_6 = self.upsampler_6(upsample_feature_5, downsample_feature_1)
+		output = torch.nn.functional.interpolate(upsample_feature_6, scale_factor = 2, mode = 'bilinear', align_corners = False)
+		return bottleneck_output, upsample_feature_1, upsample_feature_2, upsample_feature_3, upsample_feature_4, upsample_feature_5, upsample_feature_6, output