Unverified Commit be4ff9a3 authored by Nicolas Hug's avatar Nicolas Hug Committed by GitHub
Browse files

Documentation improvements for ops (#3634)



* some doc improvements

* flake8
Co-authored-by: default avatarFrancisco Massa <fvsmassa@gmail.com>
parent 27156ad9
......@@ -28,9 +28,8 @@ def nms(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:
iou_threshold (float): discards all overlapping boxes with IoU > iou_threshold
Returns:
keep (Tensor): int64 tensor with the indices
of the elements that have been kept
by NMS, sorted in decreasing order of scores
Tensor: int64 tensor with the indices of the elements that have been kept
by NMS, sorted in decreasing order of scores
"""
_assert_has_ops()
return torch.ops.torchvision.nms(boxes, scores, iou_threshold)
......@@ -57,9 +56,8 @@ def batched_nms(
iou_threshold (float): discards all overlapping boxes with IoU > iou_threshold
Returns:
keep (Tensor): int64 tensor with the indices of
the elements that have been kept by NMS, sorted
in decreasing order of scores
Tensor: int64 tensor with the indices of the elements that have been kept by NMS, sorted
in decreasing order of scores
"""
# Benchmarks that drove the following thresholds are at
# https://github.com/pytorch/vision/issues/1311#issuecomment-781329339
......@@ -117,8 +115,8 @@ def remove_small_boxes(boxes: Tensor, min_size: float) -> Tensor:
min_size (float): minimum size
Returns:
keep (Tensor[K]): indices of the boxes that have both sides
larger than min_size
Tensor[K]: indices of the boxes that have both sides
larger than min_size
"""
ws, hs = boxes[:, 2] - boxes[:, 0], boxes[:, 3] - boxes[:, 1]
keep = (ws >= min_size) & (hs >= min_size)
......@@ -136,7 +134,7 @@ def clip_boxes_to_image(boxes: Tensor, size: Tuple[int, int]) -> Tensor:
size (Tuple[height, width]): size of the image
Returns:
clipped_boxes (Tensor[N, 4])
Tensor[N, 4]: clipped boxes
"""
dim = boxes.dim()
boxes_x = boxes[..., 0::2]
......@@ -162,6 +160,7 @@ def box_convert(boxes: Tensor, in_fmt: str, out_fmt: str) -> Tensor:
Supported in_fmt and out_fmt are:
'xyxy': boxes are represented via corners, x1, y1 being top left and x2, y2 being bottom right.
This is the format that torchvision utilities expect.
'xywh' : boxes are represented via corner, width and height, x1, y2 being top left, w, h being width and height.
......@@ -174,7 +173,7 @@ def box_convert(boxes: Tensor, in_fmt: str, out_fmt: str) -> Tensor:
out_fmt (str): Output format of given boxes. Supported formats are ['xyxy', 'xywh', 'cxcywh']
Returns:
boxes (Tensor[N, 4]): Boxes into converted format.
Tensor[N, 4]: Boxes into converted format.
"""
allowed_fmts = ("xyxy", "xywh", "cxcywh")
......@@ -215,7 +214,7 @@ def _upcast(t: Tensor) -> Tensor:
def box_area(boxes: Tensor) -> Tensor:
"""
Computes the area of a set of bounding boxes, which are specified by its
Computes the area of a set of bounding boxes, which are specified by their
(x1, y1, x2, y2) coordinates.
Args:
......@@ -224,7 +223,7 @@ def box_area(boxes: Tensor) -> Tensor:
``0 <= x1 < x2`` and ``0 <= y1 < y2``.
Returns:
area (Tensor[N]): area for each box
Tensor[N]: the area for each box
"""
boxes = _upcast(boxes)
return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
......@@ -249,17 +248,17 @@ def _box_inter_union(boxes1: Tensor, boxes2: Tensor) -> Tuple[Tensor, Tensor]:
def box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
"""
Return intersection-over-union (Jaccard index) of boxes.
Return intersection-over-union (Jaccard index) between two sets of boxes.
Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with
``0 <= x1 < x2`` and ``0 <= y1 < y2``.
Args:
boxes1 (Tensor[N, 4])
boxes2 (Tensor[M, 4])
boxes1 (Tensor[N, 4]): first set of boxes
boxes2 (Tensor[M, 4]): second set of boxes
Returns:
iou (Tensor[N, M]): the NxM matrix containing the pairwise IoU values for every element in boxes1 and boxes2
Tensor[N, M]: the NxM matrix containing the pairwise IoU values for every element in boxes1 and boxes2
"""
inter, union = _box_inter_union(boxes1, boxes2)
iou = inter / union
......@@ -269,17 +268,17 @@ def box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
# Implementation adapted from https://github.com/facebookresearch/detr/blob/master/util/box_ops.py
def generalized_box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
"""
Return generalized intersection-over-union (Jaccard index) of boxes.
Return generalized intersection-over-union (Jaccard index) between two sets of boxes.
Both sets of boxes are expected to be in ``(x1, y1, x2, y2)`` format with
``0 <= x1 < x2`` and ``0 <= y1 < y2``.
Args:
boxes1 (Tensor[N, 4])
boxes2 (Tensor[M, 4])
boxes1 (Tensor[N, 4]): first set of boxes
boxes2 (Tensor[M, 4]): second set of boxes
Returns:
generalized_iou (Tensor[N, M]): the NxM matrix containing the pairwise generalized_IoU values
Tensor[N, M]: the NxM matrix containing the pairwise generalized IoU values
for every element in boxes1 and boxes2
"""
......
......@@ -44,7 +44,7 @@ def deform_conv2d(
convolution kernel. Default: None
Returns:
output (Tensor[batch_sz, out_channels, out_h, out_w]): result of convolution
Tensor[batch_sz, out_channels, out_h, out_w]: result of convolution
Examples::
......@@ -105,7 +105,7 @@ def deform_conv2d(
class DeformConv2d(nn.Module):
"""
See deform_conv2d
See :func:`deform_conv2d`.
"""
def __init__(
......
......@@ -38,7 +38,7 @@ def ps_roi_align(
ceil(roi_width / pooled_w), and likewise for height). Default: -1
Returns:
output (Tensor[K, C, output_size[0], output_size[1]])
Tensor[K, C, output_size[0], output_size[1]]: The pooled RoIs
"""
_assert_has_ops()
check_roi_boxes_shape(boxes)
......@@ -55,7 +55,7 @@ def ps_roi_align(
class PSRoIAlign(nn.Module):
"""
See ps_roi_align
See :func:`ps_roi_align`.
"""
def __init__(
self,
......
......@@ -32,7 +32,7 @@ def ps_roi_pool(
the box coordinates. Default: 1.0
Returns:
output (Tensor[K, C, output_size[0], output_size[1]])
Tensor[K, C, output_size[0], output_size[1]]: The pooled RoIs.
"""
_assert_has_ops()
check_roi_boxes_shape(boxes)
......@@ -48,7 +48,7 @@ def ps_roi_pool(
class PSRoIPool(nn.Module):
"""
See ps_roi_pool
See :func:`ps_roi_pool`.
"""
def __init__(self, output_size: int, spatial_scale: float):
super(PSRoIPool, self).__init__()
......
......@@ -42,7 +42,7 @@ def roi_align(
This version in Detectron2
Returns:
output (Tensor[K, C, output_size[0], output_size[1]])
Tensor[K, C, output_size[0], output_size[1]]: The pooled RoIs.
"""
_assert_has_ops()
check_roi_boxes_shape(boxes)
......@@ -57,7 +57,7 @@ def roi_align(
class RoIAlign(nn.Module):
"""
See roi_align
See :func:`roi_align`.
"""
def __init__(
self,
......
......@@ -32,7 +32,7 @@ def roi_pool(
the box coordinates. Default: 1.0
Returns:
output (Tensor[K, C, output_size[0], output_size[1]])
Tensor[K, C, output_size[0], output_size[1]]: The pooled RoIs.
"""
_assert_has_ops()
check_roi_boxes_shape(boxes)
......@@ -47,7 +47,7 @@ def roi_pool(
class RoIPool(nn.Module):
"""
See roi_pool
See :func:`roi_pool`.
"""
def __init__(self, output_size: BroadcastingList2[int], spatial_scale: float):
super(RoIPool, self).__init__()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment