[Typing][A-87] Add type annotations for python/paddle/vision/datasets/flowers.py (#65504)

enkilee · web-flow · commit 6d40b68bbc66 · 2024-06-28T22:46:17.000+08:00
diff --git a/python/paddle/vision/datasets/flowers.py b/python/paddle/vision/datasets/flowers.py
@@ -12,6 +12,19 @@
 # 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.vision.transforms.transforms import _Transform
+
+    from ..image import _ImageBackend, _ImageDataType
+
+    _DatasetMode = Literal["train", "valid", "test"]
+
 import os
 import tarfile
 
@@ -44,16 +57,16 @@ class Flowers(Dataset):
     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/flowers/.
-        label_file (str, optional): Path to label file, can be set None if
+        label_file (str|None, optional): Path to label file, can be set None if
             :attr:`download` is True. Default: None, default data path: ~/.cache/paddle/dataset/flowers/.
-        setid_file (str, optional): Path to subset index file, can be set
+        setid_file (str|None, optional): Path to subset index file, can be set
             None if :attr:`download` is True. Default: None, default data path: ~/.cache/paddle/dataset/flowers/.
         mode (str, optional): Either train or test mode. Default 'train'.
-        transform (Callable, 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:
+        transform (Callable|None, optional): transform to perform on image, None for no transform. Default: None.
+        download (bool|None, optional): download dataset automatically if :attr:`data_file` is None. Default: True.
+        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.
@@ -101,20 +114,27 @@ class Flowers(Dataset):
 
             >>> for img, label in itertools.islice(iter(flowers_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, 96] [1]
     """
 
+    backend: _ImageBackend
+    data_file: str | None
+    label_file: str | None
+    setid_file: str | None
+    mode: _DatasetMode
+    transform: _Transform[Any, Any] | None
+
     def __init__(
         self,
-        data_file=None,
-        label_file=None,
-        setid_file=None,
-        mode='train',
-        transform=None,
-        download=True,
-        backend=None,
-    ):
+        data_file: str | None = None,
+        label_file: str | None = None,
+        setid_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',
@@ -167,7 +187,9 @@ def __init__(
         self.labels = scio.loadmat(label_file)['labels'][0]
         self.indexes = scio.loadmat(setid_file)[flag][0]
 
-    def __getitem__(self, idx):
+    def __getitem__(
+        self, idx: int
+    ) -> tuple[_ImageDataType, npt.NDArray[np.int64]]:
         index = self.indexes[idx]
         label = np.array([self.labels[index - 1]])
         img_name = "jpg/image_%05d.jpg" % index
