unet.py

# -*- coding: utf-8 -*-
"""Unet.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1P0m65X4HFLUsx3qcZLfKKlKLgBJvTnJw
"""

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms.functional as TF

class DoubleConv(nn.Module):
  def __init__(self, in_channels, out_channels):
    super(DoubleConv, self).__init__()
    self.conv = nn.Sequential(
        nn.Conv2d(in_channels, out_channels, 3, 1, 1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True),
        nn.Conv2d(out_channels, out_channels, 3, 1, 1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True)
    )

  def forward(self, x):
    return self.conv(x)

class UNET(nn.Module):
  def __init__(self, in_channels=3, out_channels=1, features=[64, 128, 256, 512]):
    super(UNET, self).__init__()
    self.up_sampling = nn.ModuleList()
    self.down_sampling = nn.ModuleList()
    self.max_pool = nn.MaxPool2d(kernel_size=2, stride=2)

    # Down the U network
    for feature in features:
      self.down_sampling.append(DoubleConv(in_channels, feature))
      in_channels = feature

    # Up the U network
    for feature in reversed(features):
      self.up_sampling.append(nn.ConvTranspose2d(feature * 2, feature, kernel_size=2, stride=2))
      self.up_sampling.append(DoubleConv(feature * 2, feature))

    self.bottle_neck = DoubleConv(features[-1], features[-1] * 2)
    self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)

  def forward(self, x):
    skip_connections = []
    for down_sample in self.down_sampling:
      x = down_sample(x)
      skip_connections.append(x)
      x = self.max_pool(x)

    x = self.bottle_neck(x)
    skip_connections = skip_connections[::-1]

    for i in range(0, len(self.up_sampling), 2):
      x = self.up_sampling[i](x)
      skip_connection = skip_connections[i//2]

      if x.shape != skip_connection.shape:
        x = TF.resize(x, size=skip_connection.shape[2:])

      concat_skip = torch.cat((skip_connection, x), dim=1)
      x = self.up_sampling[i + 1](concat_skip)

    return self.final_conv(x)

x = torch.randn((3, 1, 164, 164))
model = UNET(in_channels=1, out_channels=1)
print(x)

preds = model(x)

print(preds)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
unet = UNET(in_channels=1, out_channels=1).to(device)

criteria = nn.CrossEntropyLoss()
optimizer = optim.Adam(unet.parameters(), lr=0.001)

from torch.utils.data import Dataset, DataLoader

class CustomImageDataset(Dataset):
  def __init__(self, num_samples, image_size):
    self.num_samples = num_samples
    self.image_size = image_size

  def __len__(self):
    return self.num_samples

  def __getitem__(self, index):
    image = torch.randn(1, self.image_size, self.image_size)
    target_mask = torch.randint(0, 2, (1, self.image_size, self.image_size)).float()
    return image, target_mask

dataset = CustomImageDataset(1000, 256)
batch_size = 6
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# Training Loop
num_epochs = 5

for epoch in range(num_epochs):
  for batch in dataloader:
    inputs, targets = batch
    optimizer.zero_grad()
    outputs = unet(inputs.to(device))
    loss = criteria(outputs, targets.to(device))
    loss.backward()
    optimizer.step()
    print(f"Epoch: {epoch}, Loss: {loss}")