def forward(self, images, targets=None):
"""
Arguments:
#images (list[Tensor] or ImageList): images to be processed
targets (list[BoxList]): ground-truth boxes present in the image (optional)
Returns:
result (list[BoxList] or dict[Tensor]): the output from the model.
During training, it returns a dict[Tensor] which contains the losses.
During testing, it returns list[BoxList] contains additional fields
like `scores`, `labels` and `mask` (for Mask R-CNN models).
"""
if self.training and targets is None:
raise ValueError("In training mode, targets should be passed")
if self.training:
images["cur"] = to_image_list(images["cur"])
images["ref"] = [to_image_list(image) for image in images["ref"]]
return self._forward_train(images["cur"], images["ref"], targets)
else:
images["cur"] = to_image_list(images["cur"])
infos = images.copy()
infos.pop("cur")
return self._forward_test(images["cur"], infos)
def __call__(self, batch):
transposed_batch = list(zip(*batch))
if self.method in ("base", ):
images = to_image_list(transposed_batch[0], self.size_divisible)
elif self.method in ("rdn", "mega", "fgfa", "dff"):
assert len(
transposed_batch[0]
) == 1, "Currently 1 gpu could only hold 1 image. Please modify SOLVER.IMS_PER_BATCH and TEST.IMS_PER_BATCH to ensure this."
images = {}
for key in transposed_batch[0][0].keys():
if key == "cur":
images["cur"] = to_image_list(
(transposed_batch[0][0]["cur"], ), self.size_divisible)
if key not in ("ref", "ref_l", "ref_m", "ref_g"):
images[key] = transposed_batch[0][0][key]
else:
if transposed_batch[0][0][key]:
images[key] = [
to_image_list((img, ), self.size_divisible)
for img in transposed_batch[0][0][key]
]
else:
images[key] = []
else:
raise NotImplementedError("method {} not supported yet.".format(
self.method))
targets = transposed_batch[1]
img_ids = transposed_batch[2]
return images, targets, img_ids
def __call__(self, batch):
transposed_batch = list(zip(*batch))
if self.method in ("base", ):
images = to_image_list(transposed_batch[0], self.size_divisible)
elif self.method in ("rdn", "mega", "fgfa", "dff"):
images = {}
for key in transposed_batch[0][0].keys():
if key == "cur":
images["cur"] = to_image_list(
(transposed_batch[0][0]["cur"], ), self.size_divisible)
if key not in ("ref", "ref_l", "ref_m", "ref_g"):
images[key] = transposed_batch[0][0][key]
else:
if transposed_batch[0][0][key]:
images[key] = [
to_image_list((img, ), self.size_divisible)
for img in transposed_batch[0][0][key]
]
else:
images[key] = []
else:
raise NotImplementedError("method {} not supported yet.".format(
self.method))
targets = transposed_batch[1]
img_ids = transposed_batch[2]
return images, targets, img_ids
def perform_transform(self, original_image):
image = self.transforms(original_image)
image_list = to_image_list(image,
self.cfg.DATALOADER.SIZE_DIVISIBILITY)
image_list = image_list.to(self.device)
return image_list
def create_random_input(cfg, device):
ret = []
for x in cfg.INPUT.MIN_SIZE_TRAIN:
ret.append(torch.rand(3, x, int(x * 1.2)))
ret = to_image_list(ret, cfg.DATALOADER.SIZE_DIVISIBILITY)
ret = ret.to(device)
return ret
def forward(self, images, targets=None):
"""
Arguments:
images (list[Tensor] or ImageList): images to be processed
targets (list[BoxList]): ground-truth boxes present in the image (optional)
Returns:
result (list[BoxList] or dict[Tensor]): the output from the model.
During training, it returns a dict[Tensor] which contains the losses.
During testing, it returns list[BoxList] contains additional fields
like `scores`, `labels` and `mask` (for Mask R-CNN models).
"""
if self.training and targets is None:
raise ValueError("In training mode, targets should be passed")
images = to_image_list(images)
features = self.backbone(images.tensors)
proposals, proposal_losses = self.rpn(images, features, targets)
if self.roi_heads:
x, result, detector_losses = self.roi_heads(
features, proposals, targets)
else:
# RPN-only models don't have roi_heads
x = features
result = proposals
detector_losses = {}
if self.training:
losses = {}
losses.update(detector_losses)
losses.update(proposal_losses)
return losses
return result
def im_detect_bbox(model, images, target_scale, target_max_size, device):
"""
Performs bbox detection on the original image.
"""
transform = TT.Compose([
T.Resize(target_scale, target_max_size),
TT.ToTensor(),
T.Normalize(mean=cfg.INPUT.PIXEL_MEAN,
std=cfg.INPUT.PIXEL_STD,
to_bgr255=cfg.INPUT.TO_BGR255)
])
images = [transform(image) for image in images]
images = to_image_list(images, cfg.DATALOADER.SIZE_DIVISIBILITY)
return model(images.to(device))
def im_detect_bbox_hflip(model, images, target_scale, target_max_size, device):
"""
Performs bbox detection on the horizontally flipped image.
Function signature is the same as for im_detect_bbox.
"""
transform = TT.Compose([
T.Resize(target_scale, target_max_size),
TT.RandomHorizontalFlip(1.0),
TT.ToTensor(),
T.Normalize(mean=cfg.INPUT.PIXEL_MEAN,
std=cfg.INPUT.PIXEL_STD,
to_bgr255=cfg.INPUT.TO_BGR255)
])
images = [transform(image) for image in images]
images = to_image_list(images, cfg.DATALOADER.SIZE_DIVISIBILITY)
boxlists = model(images.to(device))
# Invert the detections computed on the flipped image
boxlists_inv = [boxlist.transpose(0) for boxlist in boxlists]
return boxlists_inv