Some cleanups

Author: NicolasHug

test/test_datasets.py

@@ -2517,53 +2517,43 @@ def _meta_to_split_and_classification_ann(self, meta, idx):
 class StanfordCarsTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.StanfordCars
     REQUIRED_PACKAGES = ("scipy",)
-    ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(train=(True, False))
+    ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "test"))
 
-    def _inject_fake_data(self, tmpdir, config):
+    def inject_fake_data(self, tmpdir, config):
         import scipy.io as io
         from numpy.core.records import fromarrays
 
-        train = config["train"]
-        num_examples = 5
-        root_folder = tmpdir
+        num_examples = {"train": 5, "test": 7}[config["split"]]
+        num_classes = 3
+        base_folder = pathlib.Path(tmpdir) / "stanford_cars"
 
-        class_name = np.random.randint(0, 100, num_examples, dtype=np.uint8)
-        bbox_x1 = np.random.randint(0, 100, num_examples, dtype=np.uint8)
-        bbox_x2 = np.random.randint(0, 100, num_examples, dtype=np.uint8)
+        devkit = base_folder / "devkit"
+        devkit.mkdir(parents=True)
 
-        bbox_y1 = np.random.randint(0, 100, num_examples, dtype=np.uint8)
-        bb1ox_y2 = np.random.randint(0, 100, num_examples, dtype=np.uint8)
-        fname = [f"{i:5d}.jpg" for i in range(num_examples)]
+        if config["split"] == "train":
+            images_folder_name = "cars_train"
+            annotations_mat_path = str(devkit / "cars_train_annos.mat")
+        else:
+            images_folder_name = "cars_test"
+            annotations_mat_path = str(base_folder / "cars_test_annos_withlabels.mat")
 
-        rec_array = fromarrays(
-            [bbox_x1, bbox_y1, bbox_x2, bb1ox_y2, class_name, fname],
-            names=["bbox_x1", "bbox_y1", "bbox_x2", "bbox_y2", "class", "fname"],
+        datasets_utils.create_image_folder(
+            root=base_folder,
+            name=images_folder_name,
+            file_name_fn=lambda image_index: f"{image_index:5d}.jpg",
+            num_examples=num_examples,
         )
-        devkit = os.path.join(root_folder, "devkit")
-        os.makedirs(devkit)
-
-        random_class_names = ["Tesla Model S Sedan 2012"] * 196
-
-        io.savemat(os.path.join(devkit, "cars_meta.mat"), {"class_names": random_class_names})
 
-        if train:
-            datasets_utils.create_image_folder(
-                root=root_folder,
-                name="cars_train",
-                file_name_fn=lambda image_index: f"{image_index:5d}.jpg",
-                num_examples=num_examples,
-            )
-
-            io.savemat(f"{devkit}/cars_train_annos.mat", {"annotations": rec_array})
-        else:
+        classes = np.random.randint(1, num_classes + 1, num_examples, dtype=np.uint8)
+        fnames = [f"{i:5d}.jpg" for i in range(num_examples)]
+        rec_array = fromarrays(
+            [classes, fnames],
+            names=["class", "fname"],
+        )
+        io.savemat(annotations_mat_path, {"annotations": rec_array})
 
-            datasets_utils.create_image_folder(
-                root=root_folder,
-                name="cars_test",
-                file_name_fn=lambda image_index: f"{image_index:5d}.jpg",
-                num_examples=num_examples,
-            )
-            io.savemat(f"{root_folder}/cars_test_annos_withlabels.mat", {"annotations": rec_array})
+        random_class_names = ["random_name"] * num_classes
+        io.savemat(str(devkit / "cars_meta.mat"), {"class_names": random_class_names})
 
         return num_examples
 

torchvision/datasets/stanford_cars.py

@@ -1,23 +1,26 @@
-import os
-import os.path
+import pathlib
 from typing import Callable, Optional, Any, Tuple
 
 from PIL import Image
 
-from .utils import download_and_extract_archive, download_url
+from .utils import download_and_extract_archive, download_url, verify_str_arg
 from .vision import VisionDataset
 
 
 class StanfordCars(VisionDataset):
     """`Stanford Cars <https://ai.stanford.edu/~jkrause/cars/car_dataset.html>`_ Dataset
 
-    .. warning::
+    The Cars dataset contains 16,185 images of 196 classes of cars. The data is
+    split into 8,144 training images and 8,041 testing images, where each class
+    has been split roughly in a 50-50 split
+
+    .. note::
 
         This class needs `scipy <https://docs.scipy.org/doc/>`_ to load target files from `.mat` format.
 
     Args:
         root (string): Root directory of dataset
-        train (bool, optional):If True, creates dataset from training set, otherwise creates from test set
+        split (string, optional): The dataset split, supports ``"train"`` (default) or ``"test"``.
         transform (callable, optional): A function/transform that  takes in an PIL image
             and returns a transformed version. E.g, ``transforms.RandomCrop``
         target_transform (callable, optional): A function/transform that takes in the
@@ -26,30 +29,10 @@ class StanfordCars(VisionDataset):
             puts it in root directory. If dataset is already downloaded, it is not
             downloaded again."""
 
-    urls = (
-        "https://ai.stanford.edu/~jkrause/car196/cars_test.tgz",
-        "https://ai.stanford.edu/~jkrause/car196/cars_train.tgz",
-    )  # test and train image urls
-
-    md5s = (
-        "4ce7ebf6a94d07f1952d94dd34c4d501",
-        "065e5b463ae28d29e77c1b4b166cfe61",
-    )  # md5checksum for test and train data
-
-    annot_urls = (
-        "https://ai.stanford.edu/~jkrause/car196/cars_test_annos_withlabels.mat",
-        "https://ai.stanford.edu/~jkrause/cars/car_devkit.tgz",
-    )  # annotations and labels for test and train
-
-    annot_md5s = (
-        "b0a2b23655a3edd16d84508592a98d10",
-        "c3b158d763b6e2245038c8ad08e45376",
-    )  # md5 checksum for annotations
-
     def __init__(
         self,
         root: str,
-        train: bool = True,
+        split: str = "train",
         transform: Optional[Callable] = None,
         target_transform: Optional[Callable] = None,
         download: bool = False,
@@ -62,7 +45,16 @@ def __init__(
 
         super().__init__(root, transform=transform, target_transform=target_transform)
 
-        self.train = train
+        self._split = verify_str_arg(split, "split", ("train", "test"))
+        self._base_folder = pathlib.Path(root) / "stanford_cars"
+        devkit = self._base_folder / "devkit"
+
+        if self._split == "train":
+            self._annotations_mat_path = devkit / "cars_train_annos.mat"
+            self._images_base_path = self._base_folder / "cars_train"
+        else:
+            self._annotations_mat_path = self._base_folder / "cars_test_annos_withlabels.mat"
+            self._images_base_path = self._base_folder / "cars_test"
 
         if download:
             self.download()
@@ -72,22 +64,13 @@ def __init__(
 
         self._samples = [
             (
-                os.path.join(self.root, f"cars_{'train' if self.train else 'test'}", annotation["fname"]),
-                annotation["class"] - 1,
-                # Beware stanford cars target mapping  starts from 1
+                str(self._images_base_path / annotation["fname"]),
+                annotation["class"] - 1,  # Original target mapping  starts from 1, hence -1
             )
-            for annotation in sio.loadmat(
-                os.path.join(
-                    self.root,
-                    *["devkit", "cars_train_annos.mat"] if self.train else ["cars_test_annos_withlabels.mat"],
-                ),
-                squeeze_me=True,
-            )["annotations"]
+            for annotation in sio.loadmat(self._annotations_mat_path, squeeze_me=True)["annotations"]
         ]
 
-        self.classes = sio.loadmat(os.path.join(self.root, "devkit", "cars_meta.mat"), squeeze_me=True)[
-            "class_names"
-        ].tolist()
+        self.classes = sio.loadmat(str(devkit / "cars_meta.mat"), squeeze_me=True)["class_names"].tolist()
         self.class_to_idx = {cls: i for i, cls in enumerate(self.classes)}
 
     def __len__(self) -> int:
@@ -108,20 +91,31 @@ def download(self) -> None:
         if self._check_exists():
             return
 
-        download_and_extract_archive(url=self.urls[self.train], download_root=self.root, md5=self.md5s[self.train])
-        download_and_extract_archive(url=self.annot_urls[1], download_root=self.root, md5=self.annot_md5s[1])
-        if not self.train:
+        download_and_extract_archive(
+            url="https://ai.stanford.edu/~jkrause/cars/car_devkit.tgz",
+            download_root=self._base_folder,
+            md5="c3b158d763b6e2245038c8ad08e45376",
+        )
+        if self._split == "train":
+            download_and_extract_archive(
+                url="https://ai.stanford.edu/~jkrause/car196/cars_train.tgz",
+                download_root=self._base_folder,
+                md5="065e5b463ae28d29e77c1b4b166cfe61",
+            )
+        else:
+            download_and_extract_archive(
+                url="https://ai.stanford.edu/~jkrause/car196/cars_test.tgz",
+                download_root=self._base_folder,
+                md5="4ce7ebf6a94d07f1952d94dd34c4d501",
+            )
             download_url(
-                url=self.annot_urls[0],
-                root=self.root,
-                md5=self.annot_md5s[0],
+                url="https://ai.stanford.edu/~jkrause/car196/cars_test_annos_withlabels.mat",
+                root=self._base_folder,
+                md5="b0a2b23655a3edd16d84508592a98d10",
             )
 
     def _check_exists(self) -> bool:
-        return (
-            os.path.exists(os.path.join(self.root, f"cars_{'train' if self.train else 'test'}"))
-            and os.path.isdir(os.path.join(self.root, f"cars_{'train' if self.train else 'test'}"))
-            and os.path.exists(os.path.join(self.root, "devkit", "cars_meta.mat"))
-            if self.train
-            else os.path.exists(os.path.join(self.root, "cars_test_annos_withlabels.mat"))
-        )
+        if not (self._base_folder / "devkit").is_dir():
+            return False
+
+        return self._annotations_mat_path.exists() and self._images_base_path.is_dir()