Commit dff538a0 authored by novelailab's avatar novelailab

read_from_id to read from danbooru id

parent 0f8be39c
......@@ -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)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment