From c88961aee7d5c3e92de4f7b78c814c860ae3773d Mon Sep 17 00:00:00 2001
From: Falko Habel <falko.habel@gmx.de>
Date: Sun, 23 Feb 2025 22:42:36 +0100
Subject: [PATCH] convert back to images

---
 src/aiunn/upsampler.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/aiunn/upsampler.py b/src/aiunn/upsampler.py
index d14133a..e5f17ce 100644
--- a/src/aiunn/upsampler.py
+++ b/src/aiunn/upsampler.py
@@ -6,23 +6,24 @@ from config import UpsamplerConfig
 # Upsampler model that uses the configuration from the base model.
 class Upsampler(AIIA):
     def __init__(self, base_model: AIIABase):
-        # Assume that base_model.config is an instance of UpsamplerConfig.
         super().__init__(base_model.config)
         self.base_model = base_model
         self.config = UpsamplerConfig(kwargs=self.base_model.config)
-        # Create the upsample layer using values from the configuration.
-        print(self.config.upsample_scale)
         self.upsample = nn.Upsample(
             scale_factor=self.config.upsample_scale,
             mode=self.config.upsample_mode,
             align_corners=self.config.upsample_align_corners
         )
+        # Add a final conversion layer to change channels from 512 to 3.
+        self.to_rgb = nn.Conv2d(in_channels=512, out_channels=3, kernel_size=1)
 
     def forward(self, x):
         x = self.base_model(x)
         x = self.upsample(x)
+        x = self.to_rgb(x)  # Convert feature map to RGB image.
         return x
 
+
     @classmethod
     def load(cls, path: str):
         """