Spaces:
Runtime error
Runtime error
File size: 1,080 Bytes
c310e19 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
from maskrcnn_benchmark.structures.image_list import to_image_list, to_image_target_list
class BatchCollator(object):
"""
From a list of samples from the dataset,
returns the batched images and targets.
This should be passed to the DataLoader
"""
def __init__(self, size_divisible=0):
self.size_divisible = size_divisible
def __call__(self, batch):
transposed_batch = list(zip(*batch))
images = to_image_list(transposed_batch[0], self.size_divisible)
targets = transposed_batch[1]
img_ids = transposed_batch[2]
# if transposed_batch[1] is None:
# images = to_image_list(transposed_batch[0], self.size_divisible)
# targets = transposed_batch[1]
# img_ids = transposed_batch[2]
# else:
# images, targets = to_image_target_list(transposed_batch[0], self.size_divisible, transposed_batch[1])
# img_ids = transposed_batch[2]
return images, targets, img_ids
|