From 7fd07d8c70fc706f5d6ec7f3330c1e92427209a4 Mon Sep 17 00:00:00 2001 From: Francisco Massa Date: Tue, 21 May 2019 05:05:53 -0700 Subject: [PATCH] Rename mask field to masks This makes it consistent with the other models, which returns nouns in plurial --- references/detection/coco_eval.py | 2 +- torchvision/models/detection/mask_rcnn.py | 4 ++-- torchvision/models/detection/roi_heads.py | 2 +- torchvision/models/detection/transform.py | 6 +++--- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/references/detection/coco_eval.py b/references/detection/coco_eval.py index d31f97da78e..7818c8365de 100644 --- a/references/detection/coco_eval.py +++ b/references/detection/coco_eval.py @@ -101,7 +101,7 @@ def prepare_for_coco_segmentation(self, predictions): scores = prediction["scores"] labels = prediction["labels"] - masks = prediction["mask"] + masks = prediction["masks"] masks = masks > 0.5 diff --git a/torchvision/models/detection/mask_rcnn.py b/torchvision/models/detection/mask_rcnn.py index 7fb4b0445c7..9f4d62b82df 100644 --- a/torchvision/models/detection/mask_rcnn.py +++ b/torchvision/models/detection/mask_rcnn.py @@ -43,7 +43,7 @@ class MaskRCNN(FasterRCNN): 0 and H and 0 and W - labels (Tensor[N]): the predicted labels for each image - scores (Tensor[N]): the scores or each prediction - - mask (Tensor[N, H, W]): the predicted masks for each instance, in 0-1 range. In order to + - masks (Tensor[N, H, W]): the predicted masks for each instance, in 0-1 range. In order to obtain the final segmentation masks, the soft masks can be thresholded, generally with a value of 0.5 (mask >= 0.5) @@ -249,7 +249,7 @@ def maskrcnn_resnet50_fpn(pretrained=False, progress=True, ``0`` and ``H`` and ``0`` and ``W`` - labels (``Tensor[N]``): the predicted labels for each image - scores (``Tensor[N]``): the scores or each prediction - - mask (``Tensor[N, H, W]``): the predicted masks for each instance, in ``0-1`` range. In order to + - masks (``Tensor[N, H, W]``): the predicted masks for each instance, in ``0-1`` range. In order to obtain the final segmentation masks, the soft masks can be thresholded, generally with a value of 0.5 (``mask >= 0.5``) diff --git a/torchvision/models/detection/roi_heads.py b/torchvision/models/detection/roi_heads.py index 09cfb66591f..0676f7fab9b 100644 --- a/torchvision/models/detection/roi_heads.py +++ b/torchvision/models/detection/roi_heads.py @@ -573,7 +573,7 @@ def forward(self, features, proposals, image_shapes, targets=None): labels = [r["labels"] for r in result] masks_probs = maskrcnn_inference(mask_logits, labels) for mask_prob, r in zip(masks_probs, result): - r["mask"] = mask_prob + r["masks"] = mask_prob losses.update(loss_mask) diff --git a/torchvision/models/detection/transform.py b/torchvision/models/detection/transform.py index df8df8b008e..7f35c9809c6 100644 --- a/torchvision/models/detection/transform.py +++ b/torchvision/models/detection/transform.py @@ -109,10 +109,10 @@ def postprocess(self, result, image_shapes, original_image_sizes): boxes = pred["boxes"] boxes = resize_boxes(boxes, im_s, o_im_s) result[i]["boxes"] = boxes - if "mask" in pred: - masks = pred["mask"] + if "masks" in pred: + masks = pred["masks"] masks = paste_masks_in_image(masks, boxes, o_im_s) - result[i]["mask"] = masks + result[i]["masks"] = masks if "keypoints" in pred: keypoints = pred["keypoints"] keypoints = resize_keypoints(keypoints, im_s, o_im_s)