def build_dataloader(batch_size, data_dir, cfg): train_dataset = data_mapper[cfg.train_dataset["name"]]( os.path.join(data_dir, cfg.train_dataset["name"], cfg.train_dataset["root"]), os.path.join(data_dir, cfg.train_dataset["name"], cfg.train_dataset["ann_file"]), remove_images_without_annotations=True, order=["image", "boxes", "boxes_category", "info"], ) train_sampler = build_sampler(train_dataset, batch_size) train_dataloader = DataLoader( train_dataset, sampler=train_sampler, transform=T.Compose( transforms=[ T.ShortestEdgeResize(cfg.train_image_short_size, cfg.train_image_max_size), T.RandomHorizontalFlip(), T.ToMode(), ], order=["image", "boxes", "boxes_category"], ), collator=DetectionPadCollator(), num_workers=2, ) return {"train": train_dataloader}
def build_dataloader(batch_size, data_dir, cfg): train_dataset = build_dataset(data_dir, cfg) train_sampler = build_sampler(train_dataset, batch_size) train_dataloader = DataLoader( train_dataset, sampler=train_sampler, transform=T.Compose( transforms=[ T.ShortestEdgeResize( cfg.train_image_short_size, cfg.train_image_max_size, sample_style="choice", ), T.RandomHorizontalFlip(), T.ToMode(), ], order=["image", "boxes", "boxes_category"], ), collator=DetectionPadCollator(), num_workers=2, ) return train_dataloader