fix for missing sugar-crepe example

LAION-AI · Feb 2, 2024 · 16020ac · 16020ac
1 parent 5f23a76
commit 16020ac
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/clip_benchmark/datasets/sugar_crepe.py b/clip_benchmark/datasets/sugar_crepe.py
@@ -8,9 +8,11 @@ def __init__(self,  root, ann_file, transform=None):
         self.root = root
         self.ann = json.load(open(ann_file))
         self.transform = transform
+        self.idx_strings = list(self.ann.keys()) # NOTE : indices may be non-contiguous
 
     def __getitem__(self, idx):
-        data = self.ann[str(idx)]
+        idx_str = self.idx_strings[idx]
+        data = self.ann[idx_str]
         img = Image.open(os.path.join(self.root, data['filename']))
         if self.transform is not None:
             img = self.transform(img)