read_from_id to read from danbooru id

basedformer/dataset.py

@@ -74,6 +74,7 @@ class ShardedImageDataset(data.Dataset):
         self.index = self.index[global_rank::world_size]
         #override possible gil locks by making the index map an nparray
         self.index = np.array(self.index)
+        self.ids = self.index.transpose(1, 0)[2]
         self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=self.threads)
 
     def __len__(self):
@@ -107,6 +108,13 @@ class ShardedImageDataset(data.Dataset):
 
         return data, id
 
+    def read_from_id(self, id):
+        #to be used standalone
+        offset, size, _ = self.index[self.ids == id][0]
+        data = self.mmap[offset:offset+size]
+        data = decode_jpeg(data)
+        return data
+
 class ImageDatasetBuilder():
     def __init__(self, folder_path, name, dataset=True, index=True, metadata=True, threads=None):
         self.folder_path = Path(folder_path)