Skip to content

Commit 4ec38d4

Browse files
authored
Expose docs for io and ops package (#1189)
* Expose docs for io and ops package Had do modify the docstrings to use Napoleon NumPy style, because Napoleon Google Style doesn't support multiple return arguments * Add video section
1 parent 9168476 commit 4ec38d4

File tree

5 files changed

+109
-47
lines changed

5 files changed

+109
-47
lines changed

docs/source/index.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,9 @@ architectures, and common image transformations for computer vision.
99
:caption: Package Reference
1010

1111
datasets
12+
io
1213
models
14+
ops
1315
transforms
1416
utils
1517

docs/source/io.rst

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
torchvision.io
2+
==============
3+
4+
.. currentmodule:: torchvision.io
5+
6+
The :mod:`torchvision.io` package provides functions for performing IO
7+
operations. They are currently specific to reading and writing video.
8+
9+
Video
10+
-----
11+
12+
.. autofunction:: read_video
13+
14+
.. autofunction:: read_video_timestamps
15+
16+
.. autofunction:: write_video

docs/source/ops.rst

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
torchvision.ops
2+
===============
3+
4+
.. currentmodule:: torchvision.ops
5+
6+
:mod:`torchvision.ops` implements operators that are specific for Computer Vision.
7+
8+
.. note::
9+
Those operators currently do not support TorchScript.
10+
11+
12+
.. autofunction:: nms
13+
.. autofunction:: roi_align
14+
.. autofunction:: roi_pool
15+
16+
.. autoclass:: RoIAlign
17+
.. autoclass:: RoIPool

torchvision/io/video.py

Lines changed: 38 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -28,11 +28,14 @@ def write_video(filename, video_array, fps, video_codec='libx264', options=None)
2828
"""
2929
Writes a 4d tensor in [T, H, W, C] format in a video file
3030
31-
Arguments:
32-
filename (str): path where the video will be saved
33-
video_array (Tensor[T, H, W, C]): tensor containing the individual frames,
34-
as a uint8 tensor in [T, H, W, C] format
35-
fps (Number): frames per second
31+
Parameters
32+
----------
33+
filename : str
34+
path where the video will be saved
35+
video_array : Tensor[T, H, W, C]
36+
tensor containing the individual frames, as a uint8 tensor in [T, H, W, C] format
37+
fps : Number
38+
frames per second
3639
"""
3740
_check_av_available()
3841
video_array = torch.as_tensor(video_array, dtype=torch.uint8).numpy()
@@ -135,18 +138,25 @@ def read_video(filename, start_pts=0, end_pts=None):
135138
Reads a video from a file, returning both the video frames as well as
136139
the audio frames
137140
138-
Arguments:
139-
filename (str): path to the video file
140-
start_pts (int, optional): the start presentation time of the video
141-
end_pts (int, optional): the end presentation time
142-
143-
Returns:
144-
vframes (Tensor[T, H, W, C]): the `T` video frames
145-
aframes (Tensor[K, L]): the audio frames, where `K` is the number of channels
146-
and `L` is the number of points
147-
info (Dict): metadata for the video and audio. Can contain the fields
148-
- video_fps (float)
149-
- audio_fps (int)
141+
Parameters
142+
----------
143+
filename : str
144+
path to the video file
145+
start_pts : int, optional
146+
the start presentation time of the video
147+
end_pts : int, optional
148+
the end presentation time
149+
150+
Returns
151+
-------
152+
vframes : Tensor[T, H, W, C]
153+
the `T` video frames
154+
aframes : Tensor[K, L]
155+
the audio frames, where `K` is the number of channels and `L` is the
156+
number of points
157+
info : Dict
158+
metadata for the video and audio. Can contain the fields video_fps (float)
159+
and audio_fps (int)
150160
"""
151161
_check_av_available()
152162

@@ -201,13 +211,18 @@ def read_video_timestamps(filename):
201211
202212
Note that the function decodes the whole video frame-by-frame.
203213
204-
Arguments:
205-
filename (str): path to the video file
214+
Parameters
215+
----------
216+
filename : str
217+
path to the video file
218+
219+
Returns
220+
-------
221+
pts : List[int]
222+
presentation timestamps for each one of the frames in the video.
223+
video_fps : int
224+
the frame rate for the video
206225
207-
Returns:
208-
pts (List[int]): presentation timestamps for each one of the frames
209-
in the video.
210-
video_fps (int): the frame rate for the video
211226
"""
212227
_check_av_available()
213228
container = av.open(filename, metadata_errors='ignore')

torchvision/ops/boxes.py

Lines changed: 36 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -11,17 +11,23 @@ def nms(boxes, scores, iou_threshold):
1111
IoU greater than iou_threshold with another (higher scoring)
1212
box.
1313
14-
Arguments:
15-
boxes (Tensor[N, 4]): boxes to perform NMS on. They
16-
are expected to be in (x1, y1, x2, y2) format
17-
scores (Tensor[N]): scores for each one of the boxes
18-
iou_threshold (float): discards all overlapping
19-
boxes with IoU < iou_threshold
20-
21-
Returns:
22-
keep (Tensor): int64 tensor with the indices
23-
of the elements that have been kept
24-
by NMS, sorted in decreasing order of scores
14+
Parameters
15+
----------
16+
boxes : Tensor[N, 4])
17+
boxes to perform NMS on. They
18+
are expected to be in (x1, y1, x2, y2) format
19+
scores : Tensor[N]
20+
scores for each one of the boxes
21+
iou_threshold : float
22+
discards all overlapping
23+
boxes with IoU < iou_threshold
24+
25+
Returns
26+
-------
27+
keep : Tensor
28+
int64 tensor with the indices
29+
of the elements that have been kept
30+
by NMS, sorted in decreasing order of scores
2531
"""
2632
_C = _lazy_import()
2733
return _C.nms(boxes, scores, iou_threshold)
@@ -34,19 +40,25 @@ def batched_nms(boxes, scores, idxs, iou_threshold):
3440
Each index value correspond to a category, and NMS
3541
will not be applied between elements of different categories.
3642
37-
Arguments:
38-
boxes (Tensor[N, 4]): boxes where NMS will be performed. They
39-
are expected to be in (x1, y1, x2, y2) format
40-
scores (Tensor[N]): scores for each one of the boxes
41-
idxs (Tensor[N]): indices of the categories for each
42-
one of the boxes.
43-
iou_threshold (float): discards all overlapping boxes
44-
with IoU < iou_threshold
45-
46-
Returns:
47-
keep (Tensor): int64 tensor with the indices of
48-
the elements that have been kept by NMS, sorted
49-
in decreasing order of scores
43+
Parameters
44+
----------
45+
boxes : Tensor[N, 4]
46+
boxes where NMS will be performed. They
47+
are expected to be in (x1, y1, x2, y2) format
48+
scores : Tensor[N]
49+
scores for each one of the boxes
50+
idxs : Tensor[N]
51+
indices of the categories for each one of the boxes.
52+
iou_threshold : float
53+
discards all overlapping boxes
54+
with IoU < iou_threshold
55+
56+
Returns
57+
-------
58+
keep : Tensor
59+
int64 tensor with the indices of
60+
the elements that have been kept by NMS, sorted
61+
in decreasing order of scores
5062
"""
5163
if boxes.numel() == 0:
5264
return torch.empty((0,), dtype=torch.int64, device=boxes.device)

0 commit comments

Comments
 (0)