Skip to content

Commit

Permalink
Better conditional for all the kinds of bands
Browse files Browse the repository at this point in the history
Signed-off-by: João Lucas de Sousa Almeida <[email protected]>
  • Loading branch information
Joao-L-S-Almeida committed Aug 9, 2024
1 parent 52f0c12 commit 3d9b5f5
Showing 1 changed file with 12 additions and 10 deletions.
22 changes: 12 additions & 10 deletions terratorch/models/backbones/utils.py
Original file line number Diff line number Diff line change
@@ -1,31 +1,33 @@
from terratorch.datasets import HLSBands

def _are_sublists_of_int(item) -> bool:
def _are_sublists_of_int(item) -> (bool, bool):

if all([isinstance(i, list) for i in item]):
if all([isinstance(i, int) for i in sum(item, [])]):
return True
return True, True
else:
return False
raise Exception(f"It's expected sublists be [int, int], but rceived {model_bands}")
elif len(item) == 2 and type(item[0]) == type(item[1]) == int:
return False, True
else:
return False
return False, False

def _estimate_in_chans(model_bands: list[HLSBands] | list[str] | tuple[int, int] = None) -> int:

# Conditional to deal with the different possible choices for the bands
# Bands as lists of strings or enum
if all([isinstance(b, str) or isinstance(b, HLSBands) for b in model_bands]):
in_chans = len(model_bands)
is_sublist, requires_special_eval = _are_sublists_of_int(model_bands)

# Bands as intervals limited by integers
elif all([isinstance(b, int) for b in model_bands] or _are_sublists_of_int(model_bands)):
if requires_special_eval:

if _are_sublists_of_int(model_bands):
if is_sublist:
in_chans = sum([i[-1] - i[0] for i in model_bands])
else:
in_chans = model_bands[-1] - model_bands[0]
else:
raise Exception(f"Expected bands to be list(str) or [int, int] but received {model_bands}")
in_chans = len(model_bands)

return in_chans


0 comments on commit 3d9b5f5

Please sign in to comment.