Merge pull request #144 from gushiqiao/main

Fix vit calib data bug
ModelTC · Oct 17, 2024 · 128c6ce · 128c6ce
2 parents 63f3365 + 15535df
commit 128c6ce
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/llmc/data/dataset/base_dataset.py b/llmc/data/dataset/base_dataset.py
@@ -110,8 +110,10 @@ def build_calib_dataset(self):
                 for name in files:
                     if name.endswith(('.jpg', '.png', '.JPEG')):
                         img_path = os.path.join(root, name)
-                        raw_image = Image.open(img_path)
+                        raw_image = Image.open(img_path).convert('RGB')
                         self.calib_dataset.append(raw_image)
+                        if len(self.calib_dataset) == self.n_samples:
+                            return
         else:
             raise ValueError(f'Unsupported data type: {self.calib_dataset_type}')
 
@@ -262,6 +264,7 @@ def img_group_samples_wo_mask(self, samples):  # without mask
                 batch = {'pixel_values': torch.cat([sample['pixel_values']
                                                     for sample in batch], dim=0)}
                 calib_samples.append(batch)
+        return calib_samples
 
     def get_calib_dataset(self):
         samples = self.get_calib_samples()