You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
321 lines
11 KiB
321 lines
11 KiB
from collections import abc
|
|
from itertools import repeat
|
|
from numbers import Number
|
|
from typing import List
|
|
|
|
import numpy as np
|
|
|
|
from .ops import ltwh2xywh, ltwh2xyxy, resample_segments, xywh2ltwh, xywh2xyxy, xyxy2ltwh, xyxy2xywh
|
|
|
|
|
|
# From PyTorch internals
|
|
def _ntuple(n):
|
|
|
|
def parse(x):
|
|
return x if isinstance(x, abc.Iterable) else tuple(repeat(x, n))
|
|
|
|
return parse
|
|
|
|
|
|
to_4tuple = _ntuple(4)
|
|
|
|
# `xyxy` means left top and right bottom
|
|
# `xywh` means center x, center y and width, height(yolo format)
|
|
# `ltwh` means left top and width, height(coco format)
|
|
_formats = ["xyxy", "xywh", "ltwh"]
|
|
|
|
__all__ = ["Bboxes"]
|
|
|
|
|
|
class Bboxes:
|
|
"""Now only numpy is supported"""
|
|
|
|
def __init__(self, bboxes, format="xyxy") -> None:
|
|
assert format in _formats
|
|
bboxes = bboxes[None, :] if bboxes.ndim == 1 else bboxes
|
|
assert bboxes.ndim == 2
|
|
assert bboxes.shape[1] == 4
|
|
self.bboxes = bboxes
|
|
self.format = format
|
|
# self.normalized = normalized
|
|
|
|
# def convert(self, format):
|
|
# assert format in _formats
|
|
# if self.format == format:
|
|
# bboxes = self.bboxes
|
|
# elif self.format == "xyxy":
|
|
# if format == "xywh":
|
|
# bboxes = xyxy2xywh(self.bboxes)
|
|
# else:
|
|
# bboxes = xyxy2ltwh(self.bboxes)
|
|
# elif self.format == "xywh":
|
|
# if format == "xyxy":
|
|
# bboxes = xywh2xyxy(self.bboxes)
|
|
# else:
|
|
# bboxes = xywh2ltwh(self.bboxes)
|
|
# else:
|
|
# if format == "xyxy":
|
|
# bboxes = ltwh2xyxy(self.bboxes)
|
|
# else:
|
|
# bboxes = ltwh2xywh(self.bboxes)
|
|
#
|
|
# return Bboxes(bboxes, format)
|
|
|
|
def convert(self, format):
|
|
assert format in _formats
|
|
if self.format == format:
|
|
return
|
|
elif self.format == "xyxy":
|
|
bboxes = xyxy2xywh(self.bboxes) if format == "xywh" else xyxy2ltwh(self.bboxes)
|
|
elif self.format == "xywh":
|
|
bboxes = xywh2xyxy(self.bboxes) if format == "xyxy" else xywh2ltwh(self.bboxes)
|
|
else:
|
|
bboxes = ltwh2xyxy(self.bboxes) if format == "xyxy" else ltwh2xywh(self.bboxes)
|
|
self.bboxes = bboxes
|
|
self.format = format
|
|
|
|
def areas(self):
|
|
self.convert("xyxy")
|
|
return (self.bboxes[:, 2] - self.bboxes[:, 0]) * (self.bboxes[:, 3] - self.bboxes[:, 1])
|
|
|
|
# def denormalize(self, w, h):
|
|
# if not self.normalized:
|
|
# return
|
|
# assert (self.bboxes <= 1.0).all()
|
|
# self.bboxes[:, 0::2] *= w
|
|
# self.bboxes[:, 1::2] *= h
|
|
# self.normalized = False
|
|
#
|
|
# def normalize(self, w, h):
|
|
# if self.normalized:
|
|
# return
|
|
# assert (self.bboxes > 1.0).any()
|
|
# self.bboxes[:, 0::2] /= w
|
|
# self.bboxes[:, 1::2] /= h
|
|
# self.normalized = True
|
|
|
|
def mul(self, scale):
|
|
"""
|
|
Args:
|
|
scale (tuple | List | int): the scale for four coords.
|
|
"""
|
|
if isinstance(scale, Number):
|
|
scale = to_4tuple(scale)
|
|
assert isinstance(scale, (tuple, list))
|
|
assert len(scale) == 4
|
|
self.bboxes[:, 0] *= scale[0]
|
|
self.bboxes[:, 1] *= scale[1]
|
|
self.bboxes[:, 2] *= scale[2]
|
|
self.bboxes[:, 3] *= scale[3]
|
|
|
|
def add(self, offset):
|
|
"""
|
|
Args:
|
|
offset (tuple | List | int): the offset for four coords.
|
|
"""
|
|
if isinstance(offset, Number):
|
|
offset = to_4tuple(offset)
|
|
assert isinstance(offset, (tuple, list))
|
|
assert len(offset) == 4
|
|
self.bboxes[:, 0] += offset[0]
|
|
self.bboxes[:, 1] += offset[1]
|
|
self.bboxes[:, 2] += offset[2]
|
|
self.bboxes[:, 3] += offset[3]
|
|
|
|
def __len__(self):
|
|
return len(self.bboxes)
|
|
|
|
@classmethod
|
|
def concatenate(cls, boxes_list: List["Bboxes"], axis=0) -> "Bboxes":
|
|
"""
|
|
Concatenates a list of Boxes into a single Bboxes
|
|
|
|
Arguments:
|
|
boxes_list (list[Bboxes])
|
|
|
|
Returns:
|
|
Bboxes: the concatenated Boxes
|
|
"""
|
|
assert isinstance(boxes_list, (list, tuple))
|
|
if not boxes_list:
|
|
return cls(np.empty(0))
|
|
assert all(isinstance(box, Bboxes) for box in boxes_list)
|
|
|
|
if len(boxes_list) == 1:
|
|
return boxes_list[0]
|
|
return cls(np.concatenate([b.bboxes for b in boxes_list], axis=axis))
|
|
|
|
def __getitem__(self, index) -> "Bboxes":
|
|
"""
|
|
Args:
|
|
index: int, slice, or a BoolArray
|
|
|
|
Returns:
|
|
Bboxes: Create a new :class:`Bboxes` by indexing.
|
|
"""
|
|
if isinstance(index, int):
|
|
return Bboxes(self.bboxes[index].view(1, -1))
|
|
b = self.bboxes[index]
|
|
assert b.ndim == 2, f"Indexing on Bboxes with {index} failed to return a matrix!"
|
|
return Bboxes(b)
|
|
|
|
|
|
class Instances:
|
|
|
|
def __init__(self, bboxes, segments=[], keypoints=None, bbox_format="xywh", normalized=True) -> None:
|
|
"""
|
|
Args:
|
|
bboxes (ndarray): bboxes with shape [N, 4].
|
|
segments (list | ndarray): segments.
|
|
keypoints (ndarray): keypoints with shape [N, 17, 2].
|
|
"""
|
|
self._bboxes = Bboxes(bboxes=bboxes, format=bbox_format)
|
|
self.keypoints = keypoints
|
|
self.normalized = normalized
|
|
|
|
if len(segments) > 0:
|
|
# list[np.array(1000, 2)] * num_samples
|
|
segments = resample_segments(segments)
|
|
# (N, 1000, 2)
|
|
segments = np.stack(segments, axis=0)
|
|
else:
|
|
segments = np.zeros((0, 1000, 2), dtype=np.float32)
|
|
self.segments = segments
|
|
|
|
def convert_bbox(self, format):
|
|
self._bboxes.convert(format=format)
|
|
|
|
def bbox_areas(self):
|
|
self._bboxes.areas()
|
|
|
|
def scale(self, scale_w, scale_h, bbox_only=False):
|
|
"""this might be similar with denormalize func but without normalized sign"""
|
|
self._bboxes.mul(scale=(scale_w, scale_h, scale_w, scale_h))
|
|
if bbox_only:
|
|
return
|
|
self.segments[..., 0] *= scale_w
|
|
self.segments[..., 1] *= scale_h
|
|
if self.keypoints is not None:
|
|
self.keypoints[..., 0] *= scale_w
|
|
self.keypoints[..., 1] *= scale_h
|
|
|
|
def denormalize(self, w, h):
|
|
if not self.normalized:
|
|
return
|
|
self._bboxes.mul(scale=(w, h, w, h))
|
|
self.segments[..., 0] *= w
|
|
self.segments[..., 1] *= h
|
|
if self.keypoints is not None:
|
|
self.keypoints[..., 0] *= w
|
|
self.keypoints[..., 1] *= h
|
|
self.normalized = False
|
|
|
|
def normalize(self, w, h):
|
|
if self.normalized:
|
|
return
|
|
self._bboxes.mul(scale=(1 / w, 1 / h, 1 / w, 1 / h))
|
|
self.segments[..., 0] /= w
|
|
self.segments[..., 1] /= h
|
|
if self.keypoints is not None:
|
|
self.keypoints[..., 0] /= w
|
|
self.keypoints[..., 1] /= h
|
|
self.normalized = True
|
|
|
|
def add_padding(self, padw, padh):
|
|
# handle rect and mosaic situation
|
|
assert not self.normalized, "you should add padding with absolute coordinates."
|
|
self._bboxes.add(offset=(padw, padh, padw, padh))
|
|
self.segments[..., 0] += padw
|
|
self.segments[..., 1] += padh
|
|
if self.keypoints is not None:
|
|
self.keypoints[..., 0] += padw
|
|
self.keypoints[..., 1] += padh
|
|
|
|
def __getitem__(self, index) -> "Instances":
|
|
"""
|
|
Args:
|
|
index: int, slice, or a BoolArray
|
|
|
|
Returns:
|
|
Instances: Create a new :class:`Instances` by indexing.
|
|
"""
|
|
segments = self.segments[index] if len(self.segments) else self.segments
|
|
keypoints = self.keypoints[index] if self.keypoints is not None else None
|
|
bboxes = self.bboxes[index]
|
|
bbox_format = self._bboxes.format
|
|
return Instances(
|
|
bboxes=bboxes,
|
|
segments=segments,
|
|
keypoints=keypoints,
|
|
bbox_format=bbox_format,
|
|
normalized=self.normalized,
|
|
)
|
|
|
|
def flipud(self, h):
|
|
# this function may not be very logical, just for clean code when using augment flipud
|
|
self.bboxes[:, 1] = h - self.bboxes[:, 1]
|
|
self.segments[..., 1] = h - self.segments[..., 1]
|
|
if self.keypoints is not None:
|
|
self.keypoints[..., 1] = h - self.keypoints[..., 1]
|
|
|
|
def fliplr(self, w):
|
|
# this function may not be very logical, just for clean code when using augment fliplr
|
|
self.bboxes[:, 0] = w - self.bboxes[:, 0]
|
|
self.segments[..., 0] = w - self.segments[..., 0]
|
|
if self.keypoints is not None:
|
|
self.keypoints[..., 0] = w - self.keypoints[..., 0]
|
|
|
|
def clip(self, w, h):
|
|
self.convert_bbox(format="xyxy")
|
|
self.bboxes[:, [0, 2]] = self.bboxes[:, [0, 2]].clip(0, w)
|
|
self.bboxes[:, [1, 3]] = self.bboxes[:, [1, 3]].clip(0, h)
|
|
self.segments[..., 0] = self.segments[..., 0].clip(0, w)
|
|
self.segments[..., 1] = self.segments[..., 1].clip(0, h)
|
|
if self.keypoints is not None:
|
|
self.keypoints[..., 0] = self.keypoints[..., 0].clip(0, w)
|
|
self.keypoints[..., 1] = self.keypoints[..., 1].clip(0, h)
|
|
|
|
def update(self, bboxes, segments=None, keypoints=None):
|
|
new_bboxes = Bboxes(bboxes, format=self._bboxes.format)
|
|
self._bboxes = new_bboxes
|
|
if segments is not None:
|
|
self.segments = segments
|
|
if keypoints is not None:
|
|
self.keypoints = keypoints
|
|
|
|
def __len__(self):
|
|
return len(self.bboxes)
|
|
|
|
@classmethod
|
|
def concatenate(cls, instances_list: List["Instances"], axis=0) -> "Instances":
|
|
"""
|
|
Concatenates a list of Boxes into a single Bboxes
|
|
|
|
Arguments:
|
|
instances_list (list[Bboxes])
|
|
axis
|
|
|
|
Returns:
|
|
Boxes: the concatenated Boxes
|
|
"""
|
|
assert isinstance(instances_list, (list, tuple))
|
|
if not instances_list:
|
|
return cls(np.empty(0))
|
|
assert all(isinstance(instance, Instances) for instance in instances_list)
|
|
|
|
if len(instances_list) == 1:
|
|
return instances_list[0]
|
|
|
|
use_keypoint = instances_list[0].keypoints is not None
|
|
bbox_format = instances_list[0]._bboxes.format
|
|
normalized = instances_list[0].normalized
|
|
|
|
cat_boxes = np.concatenate([ins.bboxes for ins in instances_list], axis=axis)
|
|
cat_segments = np.concatenate([b.segments for b in instances_list], axis=axis)
|
|
cat_keypoints = np.concatenate([b.keypoints for b in instances_list], axis=axis) if use_keypoint else None
|
|
return cls(cat_boxes, cat_segments, cat_keypoints, bbox_format, normalized)
|
|
|
|
@property
|
|
def bboxes(self):
|
|
return self._bboxes.bboxes
|