Create a dataset from a folder of images. If you pass in a root directory it will be searched for images ending in ext (ext can be a list)
(self,
root_dir,
caption_file=None,
image_transforms=[],
ext="jpg",
default_caption="",
postprocess=None,
return_paths=False,
)
| 42 | |
| 43 | class FolderData(Dataset): |
| 44 | def __init__(self, |
| 45 | root_dir, |
| 46 | caption_file=None, |
| 47 | image_transforms=[], |
| 48 | ext="jpg", |
| 49 | default_caption="", |
| 50 | postprocess=None, |
| 51 | return_paths=False, |
| 52 | ) -> None: |
| 53 | """Create a dataset from a folder of images. |
| 54 | If you pass in a root directory it will be searched for images |
| 55 | ending in ext (ext can be a list) |
| 56 | """ |
| 57 | self.root_dir = Path(root_dir) |
| 58 | self.default_caption = default_caption |
| 59 | self.return_paths = return_paths |
| 60 | if isinstance(postprocess, DictConfig): |
| 61 | postprocess = instantiate_from_config(postprocess) |
| 62 | self.postprocess = postprocess |
| 63 | if caption_file is not None: |
| 64 | with open(caption_file, "rt") as f: |
| 65 | ext = Path(caption_file).suffix.lower() |
| 66 | if ext == ".json": |
| 67 | captions = json.load(f) |
| 68 | elif ext == ".jsonl": |
| 69 | lines = f.readlines() |
| 70 | lines = [json.loads(x) for x in lines] |
| 71 | captions = {x["file_name"]: x["text"].strip("\n") for x in lines} |
| 72 | else: |
| 73 | raise ValueError(f"Unrecognised format: {ext}") |
| 74 | self.captions = captions |
| 75 | else: |
| 76 | self.captions = None |
| 77 | |
| 78 | #print("self.captions, return paths, postprocess, default caption: ", self.captions, return_paths, postprocess, default_caption) #None False None |
| 79 | |
| 80 | if not isinstance(ext, (tuple, list, ListConfig)): |
| 81 | ext = [ext] |
| 82 | |
| 83 | # Only used if there is no caption file |
| 84 | #self.paths = [] |
| 85 | # for e in ext: |
| 86 | # self.paths.extend(sorted(list(self.root_dir.rglob(f"*img_gt.{e}")))) |
| 87 | |
| 88 | |
| 89 | img_path = '/share/phoenix/nfs05/S8/gc492/nerfw/nerfw/results/phototourism/training_data' |
| 90 | imgs = glob.glob( os.path.join(img_path, "*.png") ) # (512, 512, 3) |
| 91 | #imgs = sorted(imgs, key=extract_number) # this line extracts nfs05 instead of 000.png :( |
| 92 | imgs = [s.split("/")[-1].split('.png')[0] for s in imgs] # leave only 000, 001...2639 |
| 93 | imgs = sorted(imgs, key=lambda s: int(s)) |
| 94 | imgs = [os.path.join(img_path, i+'.png') for i in imgs] |
| 95 | self.extrinsics = np.load( os.path.join(img_path, "extrinsics.npy") ) |
| 96 | |
| 97 | # processed_images = [] |
| 98 | # processed_ext = [] |
| 99 | # skip = 11 |
| 100 | # i = 0 |
| 101 | # while i < len(imgs): |
nothing calls this directly
no test coverage detected