Machine-Learning
/
aiuNN
1
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

updated process

This commit is contained in:
Falko Victor Habel 2025-01-30 12:39:24 +01:00
parent be5bb53620
commit 4037a07764
1 changed files with 9 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
src/aiunn/finetune.py
View File

@ -23,22 +23,25 @@ class ImageDataset(Dataset):
row = self.dataframe.iloc[idx] row = self.dataframe.iloc[idx]
try: try:
# Directly use bytes data for PNG images # Convert string data to bytes if needed
low_res_bytes = row['image_512'] low_res_bytes = row['image_512'].encode() if isinstance(row['image_512'], str) else row['image_512']
high_res_bytes = row['image_1024'] 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) low_res_stream = io.BytesIO(low_res_bytes)
high_res_stream = io.BytesIO(high_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') low_res_image = Image.open(low_res_stream).convert('RGB')
high_res_image = Image.open(high_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: except Exception as e:
raise ValueError(f"Image loading failed: {str(e)}") raise ValueError(f"Image loading failed: {str(e)}")
# Apply transformations if specified
if self.transform: if self.transform:
low_res_image = self.transform(low_res_image) low_res_image = self.transform(low_res_image)
high_res_image = self.transform(high_res_image) high_res_image = self.transform(high_res_image)