Update llff.py

tne-ai · Sep 25, 2020 · e6110d9 · e6110d9
1 parent c739e49
commit e6110d9
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/datasets/llff.py b/datasets/llff.py
@@ -223,7 +223,7 @@ def read_meta(self):
                     continue
                 c2w = torch.FloatTensor(self.poses[i])
 
-                img = Image.open(image_path)
+                img = Image.open(image_path).convert('RGB')
                 assert img.size[1]*self.img_wh[0] == img.size[0]*self.img_wh[1], \
                     f'''{image_path} has different aspect ratio than img_wh, 
                         please check your data!'''
@@ -309,10 +309,10 @@ def __getitem__(self, idx):
                       'c2w': c2w}
 
             if self.split == 'val':
-                img = Image.open(self.image_path_val)
+                img = Image.open(self.image_path_val).convert('RGB')
                 img = img.resize(self.img_wh, Image.LANCZOS)
                 img = self.transform(img) # (3, h, w)
                 img = img.view(3, -1).permute(1, 0) # (h*w, 3)
                 sample['rgbs'] = img
 
-        return sample
+        return sample