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
