From 3f6e6514a9770523f248b5874c0e8a264e695a62 Mon Sep 17 00:00:00 2001
From: Falko Habel <falko.habel@gmx.de>
Date: Sun, 26 Jan 2025 22:17:37 +0100
Subject: [PATCH] rgba conversion to rgb

---
 src/aiia/data/DataLoader.py | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/src/aiia/data/DataLoader.py b/src/aiia/data/DataLoader.py
index 1f5d75a..78f4eb5 100644
--- a/src/aiia/data/DataLoader.py
+++ b/src/aiia/data/DataLoader.py
@@ -21,7 +21,13 @@ class FilePathLoader:
     def _get_image(self, item):
         try:
             path = item[self.file_path_column]
-            image = Image.open(path).convert("RGB")
+            image = Image.open(path)
+            if image.mode == 'RGBA':
+                background = Image.new('RGB', image.size, (0, 0, 0))
+                background.paste(image, mask=image.split()[3])
+                image = background
+            elif image.mode != 'RGB':
+                image = image.convert('RGB')
             return image
         except Exception as e:
             print(f"Error loading image from {path}: {e}")
@@ -69,7 +75,13 @@ class JPGImageLoader:
                 bytes_data = data
             
             img_bytes = io.BytesIO(bytes_data)
-            image = Image.open(img_bytes).convert("RGB")
+            image = Image.open(img_bytes)
+            if image.mode == 'RGBA':
+                background = Image.new('RGB', image.size, (0, 0, 0))
+                background.paste(image, mask=image.split()[3])
+                image = background
+            elif image.mode != 'RGB':
+                image = image.convert('RGB')
             return image
         except Exception as e:
             print(f"Error loading image from bytes: {e}")
@@ -167,12 +179,10 @@ class AIIADataset(torch.utils.data.Dataset):
         item = self.items[idx]
         if isinstance(item, tuple) and len(item) == 2:
             image, label = item
-            # Convert PIL image to tensor
             image = self.transform(image)
             return (image, label)
         elif isinstance(item, tuple) and len(item) == 3:
             image, task, label = item
-            # Convert PIL image to tensor first
             image = self.transform(image)
             
             if task == 'denoise':
@@ -190,7 +200,6 @@ class AIIADataset(torch.utils.data.Dataset):
                 raise ValueError(f"Unknown task: {task}")
         else:
             if isinstance(item, Image.Image):
-                # Convert single PIL image to tensor
                 return self.transform(item)
             else:
                 raise ValueError("Invalid item format.")