|
|
|
import itertools |
|
import warnings |
|
from typing import Any, Dict, List, Tuple, Union |
|
import torch |
|
|
|
|
|
class Instances: |
|
""" |
|
This class represents a list of instances in an image. |
|
It stores the attributes of instances (e.g., boxes, masks, labels, scores) as "fields". |
|
All fields must have the same ``__len__`` which is the number of instances. |
|
|
|
All other (non-field) attributes of this class are considered private: |
|
they must start with '_' and are not modifiable by a user. |
|
|
|
Some basic usage: |
|
|
|
1. Set/get/check a field: |
|
|
|
.. code-block:: python |
|
|
|
instances.gt_boxes = Boxes(...) |
|
print(instances.pred_masks) # a tensor of shape (N, H, W) |
|
print('gt_masks' in instances) |
|
|
|
2. ``len(instances)`` returns the number of instances |
|
3. Indexing: ``instances[indices]`` will apply the indexing on all the fields |
|
and returns a new :class:`Instances`. |
|
Typically, ``indices`` is a integer vector of indices, |
|
or a binary mask of length ``num_instances`` |
|
|
|
.. code-block:: python |
|
|
|
category_3_detections = instances[instances.pred_classes == 3] |
|
confident_detections = instances[instances.scores > 0.9] |
|
""" |
|
|
|
def __init__(self, image_size: Tuple[int, int], **kwargs: Any): |
|
""" |
|
Args: |
|
image_size (height, width): the spatial size of the image. |
|
kwargs: fields to add to this `Instances`. |
|
""" |
|
self._image_size = image_size |
|
self._fields: Dict[str, Any] = {} |
|
for k, v in kwargs.items(): |
|
self.set(k, v) |
|
|
|
@property |
|
def image_size(self) -> Tuple[int, int]: |
|
""" |
|
Returns: |
|
tuple: height, width |
|
""" |
|
return self._image_size |
|
|
|
def __setattr__(self, name: str, val: Any) -> None: |
|
if name.startswith("_"): |
|
super().__setattr__(name, val) |
|
else: |
|
self.set(name, val) |
|
|
|
def __getattr__(self, name: str) -> Any: |
|
if name == "_fields" or name not in self._fields: |
|
raise AttributeError("Cannot find field '{}' in the given Instances!".format(name)) |
|
return self._fields[name] |
|
|
|
def set(self, name: str, value: Any) -> None: |
|
""" |
|
Set the field named `name` to `value`. |
|
The length of `value` must be the number of instances, |
|
and must agree with other existing fields in this object. |
|
""" |
|
with warnings.catch_warnings(record=True): |
|
data_len = len(value) |
|
if len(self._fields): |
|
assert ( |
|
len(self) == data_len |
|
), "Adding a field of length {} to a Instances of length {}".format(data_len, len(self)) |
|
self._fields[name] = value |
|
|
|
def has(self, name: str) -> bool: |
|
""" |
|
Returns: |
|
bool: whether the field called `name` exists. |
|
""" |
|
return name in self._fields |
|
|
|
def remove(self, name: str) -> None: |
|
""" |
|
Remove the field called `name`. |
|
""" |
|
del self._fields[name] |
|
|
|
def get(self, name: str) -> Any: |
|
""" |
|
Returns the field called `name`. |
|
""" |
|
return self._fields[name] |
|
|
|
def get_fields(self) -> Dict[str, Any]: |
|
""" |
|
Returns: |
|
dict: a dict which maps names (str) to data of the fields |
|
|
|
Modifying the returned dict will modify this instance. |
|
""" |
|
return self._fields |
|
|
|
|
|
def to(self, *args: Any, **kwargs: Any) -> "Instances": |
|
""" |
|
Returns: |
|
Instances: all fields are called with a `to(device)`, if the field has this method. |
|
""" |
|
ret = Instances(self._image_size) |
|
for k, v in self._fields.items(): |
|
if hasattr(v, "to"): |
|
v = v.to(*args, **kwargs) |
|
ret.set(k, v) |
|
return ret |
|
|
|
def __getitem__(self, item: Union[int, slice, torch.BoolTensor]) -> "Instances": |
|
""" |
|
Args: |
|
item: an index-like object and will be used to index all the fields. |
|
|
|
Returns: |
|
If `item` is a string, return the data in the corresponding field. |
|
Otherwise, returns an `Instances` where all fields are indexed by `item`. |
|
""" |
|
if type(item) == int: |
|
if item >= len(self) or item < -len(self): |
|
raise IndexError("Instances index out of range!") |
|
else: |
|
item = slice(item, None, len(self)) |
|
|
|
ret = Instances(self._image_size) |
|
for k, v in self._fields.items(): |
|
ret.set(k, v[item]) |
|
return ret |
|
|
|
def __len__(self) -> int: |
|
for v in self._fields.values(): |
|
|
|
return v.__len__() |
|
raise NotImplementedError("Empty Instances does not support __len__!") |
|
|
|
def __iter__(self): |
|
raise NotImplementedError("`Instances` object is not iterable!") |
|
|
|
@staticmethod |
|
def cat(instance_lists: List["Instances"]) -> "Instances": |
|
""" |
|
Args: |
|
instance_lists (list[Instances]) |
|
|
|
Returns: |
|
Instances |
|
""" |
|
assert all(isinstance(i, Instances) for i in instance_lists) |
|
assert len(instance_lists) > 0 |
|
if len(instance_lists) == 1: |
|
return instance_lists[0] |
|
|
|
image_size = instance_lists[0].image_size |
|
if not isinstance(image_size, torch.Tensor): |
|
for i in instance_lists[1:]: |
|
assert i.image_size == image_size |
|
ret = Instances(image_size) |
|
for k in instance_lists[0]._fields.keys(): |
|
values = [i.get(k) for i in instance_lists] |
|
v0 = values[0] |
|
if isinstance(v0, torch.Tensor): |
|
values = torch.cat(values, dim=0) |
|
elif isinstance(v0, list): |
|
values = list(itertools.chain(*values)) |
|
elif hasattr(type(v0), "cat"): |
|
values = type(v0).cat(values) |
|
else: |
|
raise ValueError("Unsupported type {} for concatenation".format(type(v0))) |
|
ret.set(k, values) |
|
return ret |
|
|
|
def __str__(self) -> str: |
|
s = self.__class__.__name__ + "(" |
|
s += "num_instances={}, ".format(len(self)) |
|
s += "image_height={}, ".format(self._image_size[0]) |
|
s += "image_width={}, ".format(self._image_size[1]) |
|
s += "fields=[{}])".format(", ".join((f"{k}: {v}" for k, v in self._fields.items()))) |
|
return s |
|
|
|
__repr__ = __str__ |
|
|
