[Typing][A-86, A-92] Add type annotations for python/paddle/vision/{datasets/cifar.py,image.py} (#65386)

86kkd · SigureMo · web-flow · commit e26bf13918d9 · 2024-06-24T11:26:50.000+08:00
---------

Co-authored-by: Nyakku Shigure &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/vision/datasets/cifar.py b/python/paddle/vision/datasets/cifar.py
@@ -11,6 +11,20 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal
+
+if TYPE_CHECKING:
+    import numpy.typing as npt
+
+    from paddle._typing.dtype_like import _DTypeLiteral
+    from paddle.vision.transforms.transforms import _Transform
+
+    from ..image import _ImageBackend, _ImageDataType
+
+    _DatasetMode = Literal["train", "test"]
+
 
 import pickle
 import tarfile
@@ -44,12 +58,12 @@ class Cifar10(Dataset):
     dataset, which has 10 categories.
 
     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/cifar
         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.
@@ -85,7 +99,7 @@ class Cifar10(Dataset):
             ...             mean=[0.5, 0.5, 0.5],
             ...             std=[0.5, 0.5, 0.5],
             ...             to_rgb=True,
-            ...            ),
+            ...         ),
             ...     ]
             ... )
             >>> cifar10_test = Cifar10(
@@ -98,19 +112,25 @@ class Cifar10(Dataset):
 
             >>> for img, label in itertools.islice(iter(cifar10_test), 5):  # only show first 5 images
             ...     # do something with img and label
-            ...     print(type(img), img.shape, label)
+            ...     print(type(img), img.shape, label)  # type: ignore
             ...     # <class 'paddle.Tensor'> [3, 64, 64] 3
 
     """
 
+    mode: _DatasetMode
+    backend: _ImageBackend
+    data_file: str | None
+    transform: _Transform[Any, Any] | None
+    dtype: _DTypeLiteral
+
     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',
             'test',
@@ -166,7 +186,7 @@ def _load_data(self):
                 for sample, label in zip(data, labels):
                     self.data.append((sample, label))
 
-    def __getitem__(self, idx):
+    def __getitem__(self, idx: int) -> tuple[_ImageDataType, npt.NDArray[Any]]:
         image, label = self.data[idx]
         image = np.reshape(image, [3, 32, 32])
         image = image.transpose([1, 2, 0])
@@ -191,12 +211,12 @@ class Cifar100(Cifar10):
     dataset, which has 100 categories.
 
     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/cifar
         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.
@@ -246,19 +266,19 @@ class Cifar100(Cifar10):
 
             >>> for img, label in itertools.islice(iter(cifar100_test), 5):  # only show first 5 images
             ...     # do something with img and label
-            ...     print(type(img), img.shape, label)
+            ...     print(type(img), img.shape, label)  # type: ignore
             ...     # <class 'paddle.Tensor'> [3, 64, 64] 49
 
     """
 
     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:
         super().__init__(data_file, mode, transform, download, backend)
 
     def _init_url_md5_flag(self):
diff --git a/python/paddle/vision/image.py b/python/paddle/vision/image.py
@@ -12,16 +12,30 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal, Union
+
 from PIL import Image
+from typing_extensions import TypeAlias
 
 from paddle.utils import try_import
 
+if TYPE_CHECKING:
+    import numpy.typing as npt
+    from PIL.Image import Image as PILImage
+
+    from paddle import Tensor
+
+    _ImageBackend: TypeAlias = Literal["pil", "cv2", "tensor"]
+    _ImageDataType: TypeAlias = Union[Tensor, PILImage, npt.NDArray[Any]]
+
 __all__ = []
 
-_image_backend = 'pil'
+_image_backend: _ImageBackend = 'pil'
 
 
-def set_image_backend(backend):
+def set_image_backend(backend: _ImageBackend) -> None:
     """
     Specifies the backend used to load images in class :ref:`api_paddle_datasets_ImageFolder`
     and :ref:`api_paddle_datasets_DatasetFolder` . Now support backends are pillow and opencv.
@@ -61,7 +75,7 @@ def set_image_backend(backend):
 
             >>> pil_data_folder = DatasetFolder(temp_dir)
 
-            >>> for items in pil_data_folder:
+            >>> for items in pil_data_folder:  # type: ignore
             ...     break
 
             >>> print(type(items[0]))
@@ -72,7 +86,7 @@ def set_image_backend(backend):
 
             >>> cv2_data_folder = DatasetFolder(temp_dir)
 
-            >>> for items in cv2_data_folder:
+            >>> for items in cv2_data_folder:  # type: ignore
             ...     break
 
             >>> print(type(items[0]))
@@ -88,7 +102,7 @@ def set_image_backend(backend):
     _image_backend = backend
 
 
-def get_image_backend():
+def get_image_backend() -> _ImageBackend:
     """
     Gets the name of the package used to load images
 
@@ -109,7 +123,9 @@ def get_image_backend():
     return _image_backend
 
 
-def image_load(path, backend=None):
+def image_load(
+    path: str, backend: _ImageBackend | None = None
+) -> _ImageDataType | None:
     """Load an image.
 
     Args:
@@ -136,7 +152,7 @@ def image_load(path, backend=None):
 
             >>> set_image_backend('pil')
 
-            >>> pil_img = image_load(path).convert('RGB')
+            >>> pil_img = image_load(path).convert('RGB')  # type: ignore
 
             >>> print(type(pil_img))
             <class 'PIL.Image.Image'>
diff --git a/python/paddle/vision/models/vgg.py b/python/paddle/vision/models/vgg.py
@@ -138,68 +138,35 @@ def make_layers(
 
 
 cfgs = {
-    'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
+    'A': [
+        64, 'M',
+        128, 'M',
+        256, 256, 'M',
+        512, 512, 'M',
+        512, 512, 'M',
+    ],
     'B': [
-        64,
-        64,
-        'M',
-        128,
-        128,
-        'M',
-        256,
-        256,
-        'M',
-        512,
-        512,
-        'M',
-        512,
-        512,
-        'M',
+        64, 64, 'M',
+        128, 128, 'M',
+        256, 256, 'M',
+        512, 512, 'M',
+        512, 512, 'M',
     ],
     'D': [
-        64,
-        64,
-        'M',
-        128,
-        128,
-        'M',
-        256,
-        256,
-        256,
-        'M',
-        512,
-        512,
-        512,
-        'M',
-        512,
-        512,
-        512,
-        'M',
+        64, 64, 'M',
+        128, 128, 'M',
+        256, 256, 256, 'M',
+        512, 512, 512, 'M',
+        512, 512, 512, 'M',
     ],
     'E': [
-        64,
-        64,
-        'M',
-        128,
-        128,
-        'M',
-        256,
-        256,
-        256,
-        256,
-        'M',
-        512,
-        512,
-        512,
-        512,
-        'M',
-        512,
-        512,
-        512,
-        512,
-        'M',
+        64, 64, 'M',
+        128, 128, 'M',
+        256, 256, 256, 256, 'M',
+        512, 512, 512, 512, 'M',
+        512, 512, 512, 512, 'M',
     ],
-}
+}  # fmt: skip
 
 
 def _vgg(
