cat BitMasks, PolygonMasks
Summary: Add cat method to PolygonMasks and BitMasks
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import copy
+import itertools
import numpy as np
from typing import Any, Iterator, List, Union
import pycocotools.mask as mask_utils
import torch
+from detectron2.layers import cat
from detectron2.layers.roi_align import ROIAlign
from .boxes import Boxes
@@ -201,6 +203,24 @@ class BitMasks:
# not needed now
raise NotImplementedError
+ @staticmethod
+ def cat(bitmasks_list: List["BitMasks"]) -> "BitMasks":
+ """
+ Concatenates a list of BitMasks into a single BitMasks
+
+ Arguments:
+ bitmasks_list (list[BitMasks])
+
+ Returns:
+ BitMasks: the concatenated BitMasks
+ """
+ assert isinstance(bitmasks_list, (list, tuple))
+ assert len(bitmasks_list) > 0
+ assert all(isinstance(bitmask, BitMasks) for bitmask in bitmasks_list)
+
+ cat_bitmasks = type(bitmasks_list[0])(cat([bm.tensor for bm in bitmasks_list], dim=0))
+ return cat_bitmasks
+
class PolygonMasks:
"""
@@ -376,3 +396,23 @@ class PolygonMasks:
area.append(area_per_instance)
return torch.tensor(area)
+
+ @staticmethod
+ def cat(polymasks_list: List["PolygonMasks"]) -> "PolygonMasks":
+ """
+ Concatenates a list of PolygonMasks into a single PolygonMasks
+
+ Arguments:
+ polymasks_list (list[PolygonMasks])
+
+ Returns:
+ PolygonMasks: the concatenated PolygonMasks
+ """
+ assert isinstance(polymasks_list, (list, tuple))
+ assert len(polymasks_list) > 0
+ assert all(isinstance(polymask, PolygonMasks) for polymask in polymasks_list)
+
+ cat_polymasks = type(polymasks_list[0])(
+ list(itertools.chain.from_iterable(pm.polygons for pm in polymasks_list))
+ )
+ return cat_polymasks
