Extending the supported formats for bands to include list[int]

Signed-off-by: João Lucas de Sousa Almeida <[email protected]>

Author: Joao-L-S-Almeida

terratorch/datamodules/generic_pixel_wise_data_module.py

@@ -91,9 +91,9 @@ def __init__(
         test_split: Path | None = None,
         ignore_split_file_extensions: bool = True,
         allow_substring_split_file: bool = True,
-        dataset_bands: list[HLSBands | int] | None = None,
-        predict_dataset_bands: list[HLSBands | int] | None = None,
-        output_bands: list[HLSBands | int] | None = None,
+        dataset_bands: list[HLSBands | int | list[int]] | None = None,
+        predict_dataset_bands: list[HLSBands | int | list[int]] | None = None,
+        output_bands: list[HLSBands | int | list[int]] | None = None,
         constant_scale: float = 1,
         rgb_indices: list[int] | None = None,
         train_transform: A.Compose | None | list[A.BasicTransform] = None,
@@ -330,9 +330,9 @@ def __init__(
         test_split: Path | None = None,
         ignore_split_file_extensions: bool = True,
         allow_substring_split_file: bool = True,
-        dataset_bands: list[HLSBands | int] | None = None,
-        predict_dataset_bands: list[HLSBands | int] | None = None,
-        output_bands: list[HLSBands | int] | None = None,
+        dataset_bands: list[HLSBands | int | list[int]] | None = None,
+        predict_dataset_bands: list[HLSBands | int | list[int]] | None = None,
+        output_bands: list[HLSBands | int | list[int]] | None = None,
         constant_scale: float = 1,
         rgb_indices: list[int] | None = None,
         train_transform: A.Compose | None | list[A.BasicTransform] = None,

terratorch/datasets/generic_pixel_wise_dataset.py

@@ -43,8 +43,8 @@ def __init__(
         ignore_split_file_extensions: bool = True,
         allow_substring_split_file: bool = True,
         rgb_indices: list[int] | None = None,
-        dataset_bands: list[HLSBands | int] | None = None,
-        output_bands: list[HLSBands | int] | None = None,
+        dataset_bands: list[HLSBands | int | list[int]] | None = None,
+        output_bands: list[HLSBands | int | list[int]] | None = None,
         constant_scale: float = 1,
         transform: A.Compose | None = None,
         no_data_replace: float | None = None,
@@ -179,16 +179,18 @@ def _load_file(self, path, nan_replace: int | float | None = None) -> xr.DataArr
     def _generate_bands_intervals(self, bands_intervals:List[List[int]] = None):
         bands = list()
         for b_interval in bands_intervals:
-            b_interval[-1] += 1
-            bands_sublist = np.arange(*b_interval).astype(int)
+            bands_sublist = np.arange(b_interval[0], b_interval[1] + 1).astype(int).tolist()
             bands.append(bands_sublist)
         return sorted(sum(bands, []))
 
-    def _bands_defined_by_interval(self, bands_list: List[int] | List[List[int]] = None) -> bool:
+    def _bands_defined_by_interval(self, bands_list: list[int] | list[list[int]] = None) -> bool:
         if all([type(band)==int or isinstance(band, HLSBands) for band in bands_list]):
             return False
-        elif all([isinstance(band, list) for band in bands_list]):
-            return True
+        elif all([isinstance(subinterval, list) for subinterval in bands_list]):
+            if all([type(band)==int for band in sum(bands_list, [])]):
+                return True
+            else:
+                raise Exception(f"Whe using subintervals, the limits must be int.")
         else:
             raise Exception(f"Excpected List[int] or List[List[int]], but received {type(bands_list)}.")
 
@@ -206,8 +208,8 @@ def __init__(
         ignore_split_file_extensions: bool = True,
         allow_substring_split_file: bool = True,
         rgb_indices: list[str] | None = None,
-        dataset_bands: list[HLSBands | int] | None = None,
-        output_bands: list[HLSBands | int] | None = None,
+        dataset_bands: list[HLSBands | int | list[int]] | None = None,
+        output_bands: list[HLSBands | int | list[int]] | None = None,
         class_names: list[str] | None = None,
         constant_scale: float = 1,
         transform: A.Compose | None = None,