from pathlib import Path
from torch.utils.data import Dataset
import torch
import cv2

class SegmentationDataset(Dataset):
    """Loads images and masks from standard folder structure."""
    def __init__(self, images_dir, masks_dir, image_size=512, mask_suffix=''):
        self.images_dir = Path(images_dir)
        self.masks_dir = Path(masks_dir)
        self.image_size = image_size
        self.mask_suffix = mask_suffix
        self.ids = [
            f.stem for f in self.images_dir.glob('*')
            if f.suffix.lower() in ('.png', '.jpg', '.jpeg', '.tif')
        ]

    def __len__(self):
        return len(self.ids)

    def __getitem__(self, idx):
        name = self.ids[idx]

        img_path = self.images_dir / f"{name}.png"
        if not img_path.exists():
            img_path = self.images_dir / f"{name}.jpg"
        
        image = cv2.imread(str(img_path))
        if image is None:
            raise FileNotFoundError(f"Image not found or unable to read: {img_path}")
        
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        mask_path = self.masks_dir / f"{name}{self.mask_suffix}.png"
        mask = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
        if mask is None:
            raise FileNotFoundError(f"Mask not found or unable to read: {mask_path}")

        image = cv2.resize(image, (self.image_size, self.image_size))
        mask = cv2.resize(mask, (self.image_size, self.image_size),
                          interpolation=cv2.INTER_NEAREST)

        image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
        mask = torch.from_numpy(mask).long()

        return {'image': image, 'mask': mask}