tests

Author: nilsleh

tests/datasets/test_bigearthnet.py

@@ -216,7 +216,7 @@ def test_getitem_s2(self, dataset: BigEarthNetV2) -> None:
     def test_len(self, dataset: BigEarthNetV2) -> None:
         """Test dataset length."""
         if dataset.split == 'train':
-            assert len(dataset) == 2
+            assert len(dataset) == 1
         elif dataset.split == 'val':
             assert len(dataset) == 1
         else:

torchgeo/datasets/bigearthnet.py

@@ -504,7 +504,7 @@ def _download(self, url: str, filename: Path, md5: str) -> None:
             filename: output filename to write downloaded file
             md5: md5 of downloaded file
         """
-        if not os.path.exists(filename):
+        if not os.path.exists(os.path.join(self.root, filename)):
             download_url(
                 url, self.root, filename=filename, md5=md5 if self.checksum else None
             )
@@ -667,6 +667,7 @@ def __init__(
         self.download = download
         self.checksum = checksum
         self.class2idx = {c: i for i, c in enumerate(self.class_sets[num_classes])}
+        self.idx2class = {i: c for i, c in enumerate(self.class_sets[num_classes])}
         self._verify()
 
         self.metadata_df = pd.read_parquet(os.path.join(self.root, 'metadata.parquet'))
@@ -788,10 +789,8 @@ def _load_target(self, index: int) -> Tensor:
         Returns:
             the target label
         """
-        image_labels = self.metadata_df.iloc[index]['labels']
+        indices = self.metadata_df.iloc[index]['labels']
 
-        # labels -> indices
-        indices = [self.class2idx[label] for label in image_labels]
         image_target = torch.zeros(self.num_classes, dtype=torch.long)
         image_target[indices] = 1
         return image_target
@@ -839,7 +838,7 @@ def _download(self, url: str, filename: Path, md5: str) -> None:
             filename: output filename to write downloaded file
             md5: md5 of downloaded file
         """
-        if not os.path.exists(filename):
+        if not os.path.exists(os.path.join(self.root, filename)):
             download_url(
                 url, self.root, filename=filename, md5=md5 if self.checksum else None
             )