feat: expose loader argument in Country211 and EuroSAT.

Author: GdoongMathew

torchvision/datasets/country211.py

@@ -1,7 +1,7 @@
 from pathlib import Path
-from typing import Callable, Optional, Union
+from typing import Any, Callable, Optional, Union
 
-from .folder import ImageFolder
+from .folder import default_loader, ImageFolder
 from .utils import download_and_extract_archive, verify_str_arg
 
 
@@ -21,6 +21,7 @@ class Country211(ImageFolder):
  target_transform (callable, optional): A function/transform that takes in the target and transforms it.
  download (bool, optional): If True, downloads the dataset from the internet and puts it into
  ``root/country211/``. If dataset is already downloaded, it is not downloaded again.
+ loader (callable, optional): A function to load an image given its path.
  """
 
  _URL = "https://openaipublic.azureedge.net/clip/data/country211.tgz"
@@ -33,6 +34,7 @@ def __init__(
  transform: Optional[Callable] = None,
  target_transform: Optional[Callable] = None,
  download: bool = False,
+ loader: Callable[[str], Any] = default_loader,
  ) -> None:
  self._split = verify_str_arg(split, "split", ("train", "valid", "test"))
 
@@ -46,7 +48,12 @@ def __init__(
  if not self._check_exists():
  raise RuntimeError("Dataset not found. You can use download=True to download it")
 
- super().__init__(str(self._base_folder / self._split), transform=transform, target_transform=target_transform)
+ super().__init__(
+ str(self._base_folder / self._split),
+ transform=transform,
+ target_transform=target_transform,
+ loader=loader,
+ )
  self.root = str(root)
 
  def _check_exists(self) -> bool:

torchvision/datasets/eurosat.py

@@ -1,8 +1,8 @@
 import os
 from pathlib import Path
-from typing import Callable, Optional, Union
+from typing import Any, Callable, Optional, Union
 
-from .folder import ImageFolder
+from .folder import default_loader, ImageFolder
 from .utils import download_and_extract_archive
 
 
@@ -21,6 +21,7 @@ class EuroSAT(ImageFolder):
  download (bool, optional): If True, downloads the dataset from the internet and
  puts it in root directory. If dataset is already downloaded, it is not
  downloaded again. Default is False.
+ loader (callable, optional): A function to load an image given its path.
  """
 
  def __init__(
@@ -29,6 +30,7 @@ def __init__(
  transform: Optional[Callable] = None,
  target_transform: Optional[Callable] = None,
  download: bool = False,
+ loader: Callable[[str], Any] = default_loader,
  ) -> None:
  self.root = os.path.expanduser(root)
  self._base_folder = os.path.join(self.root, "eurosat")
@@ -40,7 +42,12 @@ def __init__(
  if not self._check_exists():
  raise RuntimeError("Dataset not found. You can use download=True to download it")
 
- super().__init__(self._data_folder, transform=transform, target_transform=target_transform)
+ super().__init__(
+ self._data_folder,
+ transform=transform,
+ target_transform=target_transform,
+ loader=loader,
+ )
  self.root = os.path.expanduser(root)
 
  def __len__(self) -> int: