[Typing][A-90] Add type annotations for paddle/vision/datasets/voc2012.py (#65567)

Asthestarsfalll · SigureMo · web-flow · commit d22b5d01aee5 · 2024-07-03T03:44:19.000+08:00
---------

Co-authored-by: Nyakku Shigure &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/vision/datasets/voc2012.py b/python/paddle/vision/datasets/voc2012.py
@@ -12,8 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import io
 import tarfile
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
 from PIL import Image
@@ -22,6 +25,18 @@
 from paddle.dataset.common import _check_exists_and_download
 from paddle.io import Dataset
 
+if TYPE_CHECKING:
+    import numpy.typing as npt
+
+    from paddle._typing import DTypeLike
+    from paddle.vision.transforms.transforms import _Transform
+
+    from ..image import _ImageDataType
+
+    _ImageBackend = Literal["cv2", "pil"]
+
+    _DatasetMode = Literal["train", "valid", "test"]
+
 __all__ = []
 
 VOC_URL = 'https://dataset.bj.bcebos.com/voc/VOCtrainval_11-May-2012.tar'
@@ -41,12 +56,12 @@ class VOC2012(Dataset):
     Implementation of `VOC2012 <https://host.robots.ox.ac.uk/pascal/VOC/voc2012/>`_ dataset.
 
     Args:
-        data_file (str, optional): Path to data file, can be set None if
+        data_file (str|None, optional): Path to data file, can be set None if
             :attr:`download` is True. Default: None, default data path: ~/.cache/paddle/dataset/voc2012.
         mode (str, optional): Either train or test mode. Default 'train'.
-        transform (Callable, optional): Transform to perform on image, None for no transform. Default: None.
+        transform (Callable|None, optional): Transform to perform on image, None for no transform. Default: None.
         download (bool, optional): Download dataset automatically if :attr:`data_file` is None. Default: True.
-        backend (str, optional): Specifies which type of image to be returned:
+        backend (str|None, optional): Specifies which type of image to be returned:
             PIL.Image or numpy.ndarray. Should be one of {'pil', 'cv2'}.
             If this option is not set, will get backend from :ref:`paddle.vision.get_image_backend <api_paddle_vision_get_image_backend>`,
             default backend is 'pil'. Default: None.
@@ -58,7 +73,7 @@ class VOC2012(Dataset):
 
         .. code-block:: python
 
-            >>> # doctest: +TIMEOUT(75)
+            >>> # doctest: +TIMEOUT(120)
             >>> import itertools
             >>> import paddle.vision.transforms as T
             >>> from paddle.vision.datasets import VOC2012
@@ -98,20 +113,27 @@ class VOC2012(Dataset):
 
             >>> for img, label in itertools.islice(iter(voc2012_test), 5):  # only show first 5 images
             ...     # do something with img and label
-            ...     print(type(img), img.shape)
+            ...     print(type(img), img.shape) # type: ignore
             ...     # <class 'paddle.Tensor'> [3, 281, 500]
             ...     print(type(label), label.shape)
             ...     # <class 'numpy.ndarray'> (281, 500)
     """
 
+    data_file: str | None
+    mode: _DatasetMode
+    transform: _Transform[Any, Any] | None
+    backend: _ImageBackend
+    flag: str
+    dtype: DTypeLike
+
     def __init__(
         self,
-        data_file=None,
-        mode='train',
-        transform=None,
-        download=True,
-        backend=None,
-    ):
+        data_file: str | None = None,
+        mode: _DatasetMode = 'train',
+        transform: _Transform[Any, Any] | None = None,
+        download: bool = True,
+        backend: _ImageBackend | None = None,
+    ) -> None:
         assert mode.lower() in [
             'train',
             'valid',
@@ -162,7 +184,7 @@ def _load_anno(self):
             self.data.append(data)
             self.labels.append(label)
 
-    def __getitem__(self, idx):
+    def __getitem__(self, idx: int) -> tuple[_ImageDataType, npt.NDArray[Any]]:
         data_file = self.data[idx]
         label_file = self.labels[idx]
 
@@ -183,9 +205,9 @@ def __getitem__(self, idx):
 
         return data, label
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.data)
 
-    def __del__(self):
+    def __del__(self) -> None:
         if self.data_tar:
             self.data_tar.close()
