Spaces:
Sleeping
Sleeping
File size: 2,775 Bytes
a92fb7a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | import torch
import torch.nn as nn
class DoubleConv(nn.Module):
def __init__(self, in_channels, out_channels):
"""
This module applies two consecutive convolutional layers to the input.
Args:
in_channels (int): Number of input channels.
out_channels (int): Number of output channels.
"""
super(DoubleConv, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.LeakyReLU(0.2, inplace=True),
)
def forward(self, x):
"""Forward pass of the DoubleConv module."""
x = self.conv(x)
return x
class UNet(nn.Module):
def __init__(self, in_channels=3, out_channels=3, features=[64, 128, 256, 512, 1024]):
"""
Basic U-Net model for image to image
Args:
in_channels (int, optional): Number of input channels.
out_channels (int, optional): Number of output channels.
features (list, optional): List of features for encoding and decoding.
"""
super(UNet, self).__init__()
self.ups = nn.ModuleList()
self.downs = nn.ModuleList()
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
# Encoder
for feature in features:
self.downs.append(DoubleConv(in_channels, feature))
in_channels = feature
# Decoder
for feature in reversed(features):
self.ups.append(nn.ConvTranspose2d(feature*2, feature, kernel_size=2, stride=2))
self.ups.append(DoubleConv(feature*2, feature))
self.bottleneck = DoubleConv(features[-1], features[-1]*2)
self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
def forward(self, x):
"""Forward pass of the U-Net model"""
skip_connections = []
for down in self.downs:
x = down(x)
skip_connections.append(x)
x = self.pool(x)
x = self.bottleneck(x)
skip_connections = skip_connections[::-1]
for i in range(0, len(self.ups), 2):
x = self.ups[i](x)
skip_connection = skip_connections[i//2]
if x.shape != skip_connection.shape:
x = nn.functional.interpolate(x, size=skip_connection.shape[2:], mode='bilinear', align_corners=True)
concat_skip = torch.cat((skip_connection, x), dim=1)
x = self.ups[i+1](concat_skip)
x = self.final_conv(x)
return x |