Project: s41467-025-56622-2 | Function: forward (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** UNet forward pass for image segmentation. Implements encoder-decoder architecture with skip connections for automated cytometric gating using density maps as input.

**KEY TERMS:** {
  "skip_connections": "skip connections allow the decoder to use the contextual information from the encoder",
  "downs": "encoder (the contracting path)",
  "ups": "decoder (the expansive path)",
  "bottleneck": "network bottleneck between encoder and decoder"
}

**CONTEXT CODE:** import torch
import torch.nn as nn
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, bias = False), # kernel, stride, padding, set bias to false because use batchnorm
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True),
        nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias = False), # kernel, stride, padding, set bias to false because use batchnorm
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True),
    )
  def forward(self, x):
    return self.conv(x)

class UNITO(nn.Module):
  def __init__(
    self, in_channels = 1, out_channels = 1, features = [64,128,256,512],
  ):
    super(UNITO, self).__init__()
    self.ups = nn.ModuleList()
    self.downs = nn.ModuleList()
    self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

    # Down part of UNITO
    for feature in features:
      self.downs.append(DoubleConv(in_channels, feature))
      in_channels = feature

    # Up part of UNITO
    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):
    """
    UNet forward pass for image segmentation.

    This function implements the forward pass of the UNet architecture, performing
    encoding through downsampling layers with skip connections, bottleneck processing,
    and decoding through upsampling layers. It takes input images and produces
    segmentation masks by combining features at multiple scales.

    Args:
        x: Input tensor of shape (batch_size, channels, height, width)

    Returns:
        Output tensor of shape (batch_size, out_channels, height, width)
        representing the segmentation mask
    """
    raise NotImplementedError("This function has been masked for testing purposes")

Complete the function implementation replacing "raise NotImplementedError".