Update pastis dataset

VMarsocci · Oct 15, 2024 · b7d05b9 · b7d05b9
1 parent 7ccc8db
commit b7d05b9
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 38 deletions.
diff --git a/configs/dataset/pastis.yaml b/configs/dataset/pastis.yaml
@@ -10,8 +10,8 @@ multi_modal: True
 #limited_label: False
 
 # classes
-ignore_index: 0
-num_classes: 19
+ignore_index: 19
+num_classes: 20
 classes: 
   - Background
   - Meadow
@@ -32,7 +32,7 @@ classes:
   - Orchard
   - Mixed Cereal
   - Sorghum
-  #- Void Label
+  - Void Label
 distribution:
   - 0.00000
   - 0.25675
@@ -53,7 +53,7 @@ distribution:
   - 0.02460
   - 0.00696
   - 0.00580
-  #- 0.29476
+  - 0.29476
 
 bands:
   optical:

diff --git a/pangaea/datasets/pastis.py b/pangaea/datasets/pastis.py
@@ -141,31 +141,15 @@ def __init__(
             folds = [4]
         else:
             folds = [5]
-
-        self.dataset_name = dataset_name
-        self.bands = bands
-        self.split = split
-        self.path = root_path
-        self.data_mean = data_mean
-        self.data_std = data_std
-        self.data_min = data_min
-        self.data_max = data_max
-        self.classes = classes
-        self.img_size = img_size
-        self.distribution = distribution
-
-        self.num_classes = num_classes
-        self.ignore_index = ignore_index
-        self.grid_size = multi_temporal
-        self.download_url = download_url
-        self.auto_download = auto_download
         self.modalities = ["s2", "aerial", "s1-asc"]
         self.nb_split = 1
 
         reference_date = "2018-09-01"
         self.reference_date = datetime(*map(int, reference_date.split("-")))
 
-        self.meta_patch = gpd.read_file(os.path.join(self.path, "metadata.geojson"))
+        self.meta_patch = gpd.read_file(
+            os.path.join(self.root_path, "metadata.geojson")
+        )
 
         self.num_classes = 20
 
@@ -193,19 +177,16 @@ def __getitem__(self, i: int) -> dict[str, torch.Tensor | dict[str, torch.Tensor
         part = i % (self.nb_split * self.nb_split)
         label = torch.from_numpy(
             np.load(
-                os.path.join(self.path, "ANNOTATIONS/TARGET_" + str(name) + ".npy")
+                os.path.join(self.root_path, "ANNOTATIONS/TARGET_" + str(name) + ".npy")
             )[0].astype(np.int32)
         )
-        # remove void class
-        label[label == 19] = self.ignore_index
-        # label = label[1:-1]  # remove Background and Void classes
         output = {"label": label, "name": name}
 
         for modality in self.modalities:
             if modality == "aerial":
                 with rasterio.open(
                     os.path.join(
-                        self.path,
+                        self.root_path,
                         "DATA_SPOT/PASTIS_SPOT6_RVB_1M00_2019/SPOT6_RVB_1M00_2019_"
                         + str(name)
                         + ".tif",
@@ -220,7 +201,7 @@ def __getitem__(self, i: int) -> dict[str, torch.Tensor | dict[str, torch.Tensor
                     torch.from_numpy(
                         np.load(
                             os.path.join(
-                                self.path,
+                                self.root_path,
                                 "DATA_{}".format(modality_name.upper()),
                                 "{}_{}.npy".format(modality_name.upper(), name),
                             )
@@ -237,7 +218,7 @@ def __getitem__(self, i: int) -> dict[str, torch.Tensor | dict[str, torch.Tensor
                     torch.from_numpy(
                         np.load(
                             os.path.join(
-                                self.path,
+                                self.root_path,
                                 "DATA_{}".format(modality_name.upper()),
                                 "{}_{}.npy".format(modality_name.upper(), name),
                             )
@@ -254,7 +235,7 @@ def __getitem__(self, i: int) -> dict[str, torch.Tensor | dict[str, torch.Tensor
                     torch.from_numpy(
                         np.load(
                             os.path.join(
-                                self.path,
+                                self.root_path,
                                 "DATA_{}".format(modality_name.upper()),
                                 "{}_{}.npy".format(modality_name.upper(), name),
                             )
@@ -286,7 +267,7 @@ def __getitem__(self, i: int) -> dict[str, torch.Tensor | dict[str, torch.Tensor
                     torch.from_numpy(
                         np.load(
                             os.path.join(
-                                self.path,
+                                self.root_path,
                                 "DATA_{}".format(modality_name.upper()),
                                 "{}_{}.npy".format(modality_name.upper(), name),
                             )
@@ -319,7 +300,7 @@ def __getitem__(self, i: int) -> dict[str, torch.Tensor | dict[str, torch.Tensor
                         torch.from_numpy(
                             np.load(
                                 os.path.join(
-                                    self.path,
+                                    self.root_path,
                                     "DATA_{}".format(modality_name.upper()),
                                     "{}_{}.npy".format(modality_name.upper(), name),
                                 )
@@ -337,7 +318,7 @@ def __getitem__(self, i: int) -> dict[str, torch.Tensor | dict[str, torch.Tensor
                         torch.from_numpy(
                             np.load(
                                 os.path.join(
-                                    self.path,
+                                    self.root_path,
                                     "DATA_{}".format(modality.upper()),
                                     "{}_{}.npy".format(modality.upper(), name),
                                 )
@@ -360,17 +341,17 @@ def __getitem__(self, i: int) -> dict[str, torch.Tensor | dict[str, torch.Tensor
         optical_ts = rearrange(output["s2"], "t c h w -> c t h w")
         sar_ts = rearrange(output["s1-asc"], "t c h w -> c t h w")
 
-        if self.grid_size == 1:
+        if self.multi_temporal == 1:
             # we only take the last frame
             optical_ts = optical_ts[:, -1]
             sar_ts = sar_ts[:, -1]
         else:
             # select evenly spaced samples
             optical_indexes = torch.linspace(
-                0, optical_ts.shape[1] - 1, self.grid_size, dtype=torch.long
+                0, optical_ts.shape[1] - 1, self.multi_temporal, dtype=torch.long
             )
             sar_indexes = torch.linspace(
-                0, sar_ts.shape[1] - 1, self.grid_size, dtype=torch.long
+                0, sar_ts.shape[1] - 1, self.multi_temporal, dtype=torch.long
             )
 
             optical_ts = optical_ts[:, optical_indexes]
@@ -381,7 +362,7 @@ def __getitem__(self, i: int) -> dict[str, torch.Tensor | dict[str, torch.Tensor
                 "optical": optical_ts.to(torch.float32),
                 "sar": sar_ts.to(torch.float32),
             },
-            "target": output["label"],
+            "target": output["label"].to(torch.int64),
             "metadata": {},
         }