From 4037a0776451695305ec8cba37059fef7fd168cc Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Thu, 30 Jan 2025 12:39:24 +0100 Subject: [PATCH] updated process --- src/aiunn/finetune.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/src/aiunn/finetune.py b/src/aiunn/finetune.py index 0d49d85..05fde00 100644 --- a/src/aiunn/finetune.py +++ b/src/aiunn/finetune.py @@ -23,22 +23,25 @@ class ImageDataset(Dataset): row = self.dataframe.iloc[idx] try: - # Directly use bytes data for PNG images - low_res_bytes = row['image_512'] - high_res_bytes = row['image_1024'] + # Convert string data to bytes if needed + low_res_bytes = row['image_512'].encode() if isinstance(row['image_512'], str) else row['image_512'] + high_res_bytes = row['image_1024'].encode() if isinstance(row['image_1024'], str) else row['image_1024'] - # Create in-memory streams + # Create BytesIO objects low_res_stream = io.BytesIO(low_res_bytes) high_res_stream = io.BytesIO(high_res_bytes) - # Open images with explicit RGB conversion + # Open images low_res_image = Image.open(low_res_stream).convert('RGB') high_res_image = Image.open(high_res_stream).convert('RGB') + # Close the streams + low_res_stream.close() + high_res_stream.close() + except Exception as e: raise ValueError(f"Image loading failed: {str(e)}") - # Apply transformations if specified if self.transform: low_res_image = self.transform(low_res_image) high_res_image = self.transform(high_res_image)