|
|
|
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__
|
|
|
