classBoxes:"""
This structure stores a list of boxes as a Nx4 torch.Tensor.
It supports some common methods about boxes
(`area`, `clip`, `nonempty`, etc),
and also behaves like a Tensor
(support indexing, `to(device)`, `.device`, and iteration over all boxes)
Attributes:
tensor (torch.Tensor): float matrix of Nx4.
"""
BoxSizeType = Union[List[int], Tuple[int,int]]def__init__(self, tensor: torch.Tensor):"""
Args:
tensor (Tensor[float]): a Nx4 matrix. Each row is (x1, y1, x2, y2).
"""
device = tensor.device ifisinstance(tensor, torch.Tensor)else torch.device("cpu")
tensor = torch.as_tensor(tensor, dtype=torch.float32, device=device)if tensor.numel()==0:
tensor = torch.zeros(0,4, dtype=torch.float32, device=device)assert tensor.dim()==2and tensor.size(-1)==4, tensor.size()
self.tensor = tensor
clone , 复制Boxes(新的内存空间)
defclone(self)->"Boxes":"""
Clone the Boxes.
Returns:
Boxes
"""return Boxes(self.tensor.clone())
defarea(self)-> torch.Tensor:"""
Computes the area of all the boxes.
Returns:
torch.Tensor: a vector with areas of each box.
"""
box = self.tensor
area =(box[:,2]- box[:,0])*(box[:,3]- box[:,1])return area
clip, 将框坐标clip到图片区域
defclip(self, box_size: BoxSizeType)->None:"""
Clip (in place) the boxes by limiting x coordinates to the range [0, width]
and y coordinates to the range [0, height].
Args:
box_size (height, width): The clipping box's size.
"""assert torch.isfinite(self.tensor).all(),"Box tensor contains infinite or NaN!"
h, w = box_size
self.tensor[:,0].clamp_(min=0,max=w)
self.tensor[:,1].clamp_(min=0,max=h)
self.tensor[:,2].clamp_(min=0,max=w)
self.tensor[:,3].clamp_(min=0,max=h)
nonempt, 清除小于指定大小的框
defnonempty(self, threshold:int=0)-> torch.Tensor:"""
Find boxes that are non-empty.
A box is considered empty, if either of its side is no larger than threshold.
Returns:
Tensor:
a binary vector which represents whether each box is empty
(False) or non-empty (True).
"""
box = self.tensor
widths = box[:,2]- box[:,0]
heights = box[:,3]- box[:,1]
keep =(widths > threshold)&(heights > threshold)return keep
_
_
g
e
t
_
i
t
e
m
_
_
\_\_get\_item\_\_
__get_item__,索引获取框。
"""
Returns:
Boxes: Create a new :class:`Boxes` by indexing.
The following usage are allowed:
1. `new_boxes = boxes[3]`: return a `Boxes` which contains only one box.
2. `new_boxes = boxes[2:10]`: return a slice of boxes.
3. `new_boxes = boxes[vector]`, where vector is a torch.BoolTensor
with `length = len(boxes)`. Nonzero elements in the vector will be selected.
Note that the returned Boxes might share storage with this Boxes,
subject to Pytorch's indexing semantics.
"""ifisinstance(item,int):return Boxes(self.tensor[item].view(1,-1))
b = self.tensor[item]assert b.dim()==2,"Indexing on Boxes with {} failed to return a matrix!".format(item)return Boxes(b)
definside_box(self, box_size: BoxSizeType, boundary_threshold:int=0)-> torch.Tensor:"""
Args:
box_size (height, width): Size of the reference box.
boundary_threshold (int): Boxes that extend beyond the reference box
boundary by more than boundary_threshold are considered "outside".
Returns:
a binary vector, indicating whether each box is inside the reference box.
"""
height, width = box_size
inds_inside =((self.tensor[...,0]>=-boundary_threshold)&(self.tensor[...,1]>=-boundary_threshold)&(self.tensor[...,2]< width + boundary_threshold)&(self.tensor[...,3]< height + boundary_threshold))return inds_inside
get_centers,计算框中心
defget_centers(self)-> torch.Tensor:"""
Returns:
The box centers in a Nx2 array of (x, y).
"""return(self.tensor[:,:2]+ self.tensor[:,2:])/2
scale, 归一化坐标
defscale(self, scale_x:float, scale_y:float)->None:"""
Scale the box with horizontal and vertical scaling factors
"""
self.tensor[:,0::2]*= scale_x
self.tensor[:,1::2]*= scale_y
cat, 将多个Boxes 拼接成一个
@staticmethoddefcat(boxes_list: List["Boxes"])->"Boxes":"""
Concatenates a list of Boxes into a single Boxes
Arguments:
boxes_list (list[Boxes])
Returns:
Boxes: the concatenated Boxes
"""assertisinstance(boxes_list,(list,tuple))assertlen(boxes_list)>0assertall(isinstance(box, Boxes)for box in boxes_list)
cat_boxes =type(boxes_list[0])(cat([b.tensor for b in boxes_list], dim=0))return cat_boxes