import numpy as np
from data_generator import RockGlacierDataGenerator
import traceback

gen = RockGlacierDataGenerator('datacollection/output/dataset/train/features', 'datacollection/output/dataset/train/masks', batch_size=8, target_size=(128, 128), augment=False)
try:
    feature, mask = gen._load_pair(gen.feature_files[0], gen.mask_files[0])
    print(f"File 0 feature shape: {feature.shape}, mask shape: {mask.shape}")
except Exception as e:
    print(e)