s444415/donut
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

rgb fix

Browse Source
This commit is contained in:
Michał Kozłowski 2023-01-05 14:50:49 +01:00
parent b3617532f8
commit d37f195ea6
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
utils/donut_dataset.py
View File

@ -130,8 +130,12 @@ class DonutDataset(Dataset):
""" """
sample = self.dataset[idx] sample = self.dataset[idx]
# change if not 3 channels
if sample['image'].mode != "RGB":
sample['image'] = sample['image'].convert("RGB")
# inputs # inputs
pixel_values = self.processor(sample["image"].convert("RGB"), random_padding=self.split == "train", return_tensors="pt").pixel_values pixel_values = self.processor(sample["image"], random_padding=self.split == "train", return_tensors="pt").pixel_values
pixel_values = pixel_values.squeeze() pixel_values = pixel_values.squeeze()
# targets # targets