Return labels for FER2013 if possible

torchvision/datasets/fer2013.py

@@ -25,6 +25,10 @@ class FER2013(VisionDataset):
     _RESOURCES = {
         "train": ("train.csv", "3f0dfb3d3fd99c811a1299cb947e3131"),
         "test": ("test.csv", "b02c2298636a634e8c2faabbf3ea9a23"),
+        # This one also contains both train and tests instances, and unlike test.csv it contains the labels
+        # for the test instances.
+        # It is used if it exists, otherwise "train" and "test" are used for BC, as support for "icml" was added later.
+        "icml": ("icml_face_data.csv", "b114b9e04e6949e5fe8b6a98b3892b1d"),
     }
 
     def __init__(
@@ -34,11 +38,12 @@ def __init__(
         transform: Optional[Callable] = None,
         target_transform: Optional[Callable] = None,
     ) -> None:
-        self._split = verify_str_arg(split, "split", self._RESOURCES.keys())
+        self._split = verify_str_arg(split, "split", ("train", "test"))
         super().__init__(root, transform=transform, target_transform=target_transform)
 
         base_folder = pathlib.Path(self.root) / "fer2013"
-        file_name, md5 = self._RESOURCES[self._split]
+        use_icml = (base_folder / self._RESOURCES["icml"][0]).exists()
+        file_name, md5 = self._RESOURCES["icml" if use_icml else self._split]
         data_file = base_folder / file_name
         if not check_integrity(str(data_file), md5=md5):
             raise RuntimeError(
@@ -47,14 +52,25 @@ def __init__(
                 f"https://www.kaggle.com/c/challenges-in-representation-learning-facial-expression-recognition-challenge"
             )
 
+        pixels_key = " pixels" if use_icml else "pixels"  # yes, for real
+
+        def get_img(row):
+            return torch.tensor([int(idx) for idx in row[pixels_key].split()], dtype=torch.uint8).reshape(48, 48)
+
+        def get_label(row):
+            if use_icml or self._split == "train":
+                return int(row["emotion"])
+            else:
+                return None
+
         with open(data_file, "r", newline="") as file:
-            self._samples = [
-                (
-                    torch.tensor([int(idx) for idx in row["pixels"].split()], dtype=torch.uint8).reshape(48, 48),
-                    int(row["emotion"]) if "emotion" in row else None,
-                )
-                for row in csv.DictReader(file)
-            ]
+            rows = (row for row in csv.DictReader(file))
+
+            if use_icml:
+                valid_keys = ("Training",) if self._split == "train" else ("PublicTest", "PrivateTest")
+                rows = (row for row in rows if row[" Usage"] in valid_keys)
+
+            self._samples = [(get_img(row), get_label(row)) for row in rows]
 
     def __len__(self) -> int:
         return len(self._samples)