From 413e4d898c73297902b9308f0170dcb595d3145d Mon Sep 17 00:00:00 2001 From: xander Date: Sat, 24 Aug 2024 13:36:25 +0200 Subject: [PATCH] fix img loading with transparency --- trainer/utils/io.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/trainer/utils/io.py b/trainer/utils/io.py index 9d8d5a7..86e82f2 100755 --- a/trainer/utils/io.py +++ b/trainer/utils/io.py @@ -290,6 +290,16 @@ def load_image_with_orientation(path, mode="RGB"): elif orientation == 8: image = image.rotate(90, expand=True) + if image.mode == 'P': + image = image.convert('RGBA') + + # Remove alpha channel if present + if image.mode in ('RGBA', 'LA'): + background = Image.new('RGB', image.size, (255, 255, 255)) + background.paste(image, mask=image.split()[3]) # 3 is the alpha channel + image = background + + # Convert to the desired mode return image.convert(mode)