import os
from torch.utils.data import DataLoader
from mindglow.seg2d.dataset.segmentationdataset import SegmentationDataset

def create_dataloaders(args):
    """Create train and validation dataloaders."""
    train_dataset = SegmentationDataset(
        images_dir=os.path.join(args.dataset, 'train', 'images'),
        masks_dir=os.path.join(args.dataset, 'train', 'masks'),
        image_size=args.image_size,
    )
    val_dataset = SegmentationDataset(
        images_dir=os.path.join(args.dataset, 'val', 'images'),
        masks_dir=os.path.join(args.dataset, 'val', 'masks'),
        image_size=args.image_size,
    )

    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        pin_memory=True,
        drop_last=True,
    )
    val_loader = DataLoader(
        val_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.num_workers,
        pin_memory=True,
    )
    return train_loader, val_loader