Spaces:
Sleeping
Sleeping
File size: 3,929 Bytes
749745d |
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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import torch
from maskrcnn_benchmark.structures.image_list import to_image_list
import pdb
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]
positive_map = None
positive_map_eval = None
greenlight_map = None
if isinstance(targets[0], dict):
return images, targets, img_ids, positive_map, positive_map_eval
if "greenlight_map" in transposed_batch[1][0].fields():
greenlight_map = torch.stack([i.get_field("greenlight_map") for i in transposed_batch[1]], dim=0)
if "positive_map" in transposed_batch[1][0].fields():
# we batch the positive maps here
# Since in general each batch element will have a different number of boxes,
# we collapse a single batch dimension to avoid padding. This is sufficient for our purposes.
max_len = max([v.get_field("positive_map").shape[1] for v in transposed_batch[1]])
nb_boxes = sum([v.get_field("positive_map").shape[0] for v in transposed_batch[1]])
batched_pos_map = torch.zeros((nb_boxes, max_len), dtype=torch.bool)
cur_count = 0
for v in transposed_batch[1]:
cur_pos = v.get_field("positive_map")
batched_pos_map[cur_count : cur_count + len(cur_pos), : cur_pos.shape[1]] = cur_pos
cur_count += len(cur_pos)
assert cur_count == len(batched_pos_map)
positive_map = batched_pos_map.float()
if "positive_map_eval" in transposed_batch[1][0].fields():
# we batch the positive maps here
# Since in general each batch element will have a different number of boxes,
# we collapse a single batch dimension to avoid padding. This is sufficient for our purposes.
max_len = max([v.get_field("positive_map_eval").shape[1] for v in transposed_batch[1]])
nb_boxes = sum([v.get_field("positive_map_eval").shape[0] for v in transposed_batch[1]])
batched_pos_map = torch.zeros((nb_boxes, max_len), dtype=torch.bool)
cur_count = 0
for v in transposed_batch[1]:
cur_pos = v.get_field("positive_map_eval")
batched_pos_map[cur_count : cur_count + len(cur_pos), : cur_pos.shape[1]] = cur_pos
cur_count += len(cur_pos)
assert cur_count == len(batched_pos_map)
# assert batched_pos_map.sum().item() == sum([v["positive_map"].sum().item() for v in batch[1]])
positive_map_eval = batched_pos_map.float()
return images, targets, img_ids, positive_map, positive_map_eval, greenlight_map
class BBoxAugCollator(object):
"""
From a list of samples from the dataset,
returns the images and targets.
Images should be converted to batched images in `im_detect_bbox_aug`
"""
def __call__(self, batch):
# return list(zip(*batch))
transposed_batch = list(zip(*batch))
images = transposed_batch[0]
targets = transposed_batch[1]
img_ids = transposed_batch[2]
positive_map = None
positive_map_eval = None
if isinstance(targets[0], dict):
return images, targets, img_ids, positive_map, positive_map_eval
return images, targets, img_ids, positive_map, positive_map_eval
|