release code

Dassl.ProGrad.pytorch/dassl/data/datasets/ssl/stl10.py

@@ -0,0 +1,87 @@
+import numpy as np
+import os.path as osp
+
+from dassl.utils import listdir_nohidden
+
+from ..build import DATASET_REGISTRY
+from ..base_dataset import Datum, DatasetBase
+
+
+@DATASET_REGISTRY.register()
+class STL10(DatasetBase):
+    """STL-10 dataset.
+
+    Description:
+    - 10 classes: airplane, bird, car, cat, deer, dog, horse,
+    monkey, ship, truck.
+    - Images are 96x96 pixels, color.
+    - 500 training images per class, 800 test images per class.
+    - 100,000 unlabeled images for unsupervised learning.
+
+    Reference:
+        - Coates et al. An Analysis of Single Layer Networks in
+        Unsupervised Feature Learning. AISTATS 2011.
+    """
+
+    dataset_dir = "stl10"
+
+    def __init__(self, cfg):
+        root = osp.abspath(osp.expanduser(cfg.DATASET.ROOT))
+        self.dataset_dir = osp.join(root, self.dataset_dir)
+        train_dir = osp.join(self.dataset_dir, "train")
+        test_dir = osp.join(self.dataset_dir, "test")
+        unlabeled_dir = osp.join(self.dataset_dir, "unlabeled")
+        fold_file = osp.join(
+            self.dataset_dir, "stl10_binary", "fold_indices.txt"
+        )
+
+        # Only use the first five splits
+        assert 0 <= cfg.DATASET.STL10_FOLD <= 4
+
+        train_x = self._read_data_train(
+            train_dir, cfg.DATASET.STL10_FOLD, fold_file
+        )
+        train_u = self._read_data_all(unlabeled_dir)
+        test = self._read_data_all(test_dir)
+
+        if cfg.DATASET.ALL_AS_UNLABELED:
+            train_u = train_u + train_x
+
+        super().__init__(train_x=train_x, train_u=train_u, test=test)
+
+    def _read_data_train(self, data_dir, fold, fold_file):
+        imnames = listdir_nohidden(data_dir)
+        imnames.sort()
+        items = []
+
+        list_idx = list(range(len(imnames)))
+        if fold >= 0:
+            with open(fold_file, "r") as f:
+                str_idx = f.read().splitlines()[fold]
+                list_idx = np.fromstring(str_idx, dtype=np.uint8, sep=" ")
+
+        for i in list_idx:
+            imname = imnames[i]
+            impath = osp.join(data_dir, imname)
+            label = osp.splitext(imname)[0].split("_")[1]
+            label = int(label)
+            item = Datum(impath=impath, label=label)
+            items.append(item)
+
+        return items
+
+    def _read_data_all(self, data_dir):
+        imnames = listdir_nohidden(data_dir)
+        items = []
+
+        for imname in imnames:
+            impath = osp.join(data_dir, imname)
+            label = osp.splitext(imname)[0].split("_")[1]
+            if label == "none":
+                label = -1
+            else:
+                label = int(label)
+            item = Datum(impath=impath, label=label)
+            items.append(item)
+
+        return items