"""
Custom INbreast dataset with TIF segmentation masks.
Extends the INbreast structure with pre-computed TIF masks in an ``AllMasks/``
directory. Falls back to XML-generated masks when TIF is unavailable.
"""
from pathlib import Path
from typing import Optional, Callable, Dict, Any, Tuple, List
import torch
from medical_image.utils.logging import logger
from medical_image.data.dicom_image import DicomImage
from medical_image.datasets.base_dataset import BaseDataset
from medical_image.datasets.inbreast import INbreastDataset
from medical_image.utils.mask_utils import load_tif_mask, xml_to_binary_mask
from medical_image.utils.pairing import pair_custom_inbreast, CustomINbreastSample
[docs]
class CustomINbreastDataset(BaseDataset):
"""
PyTorch Dataset for a custom INbreast-based dataset with TIF masks.
Directory structure::
root_dir/
├── AllMasks/
│ ├── 20586934_mask.tif
│ └── ...
└── INbreast Release 1.0/
├── AllDICOMs/
├── AllXML/
└── AllROI/
Each sample returns:
.. code-block:: python
{
"image": Tensor[1, H, W],
"mask": Tensor[1, H, W], # from TIF or XML fallback
"metadata": {"case_id": str, "mask_source": "tif"|"xml"|"empty", ...}
}
Args:
root_dir: Path to the custom dataset root containing ``AllMasks/``
and ``INbreast Release 1.0/``.
transform: Optional image transform.
target_transform: Optional mask transform.
target_size: Optional (H, W) resize target.
point_radius: Radius for XML single-point ROI rendering.
"""
[docs]
def __init__(
self,
root_dir: str,
transform: Optional[Callable] = None,
target_transform: Optional[Callable] = None,
target_size: Optional[Tuple[int, int]] = None,
point_radius: int = 3,
):
self.point_radius = point_radius
self._csv_metadata: Dict[str, Dict[str, Any]] = {}
super().__init__(root_dir, transform, target_transform, target_size)
# ------------------------------------------------------------------
# Build sample list
# ------------------------------------------------------------------
def _build_sample_list(self) -> None:
root = Path(self.root_dir)
# Locate directories
masks_dir = root / "AllMasks"
inbreast_dir = root / "INbreast Release 1.0"
if not inbreast_dir.exists():
raise FileNotFoundError(f"'INbreast Release 1.0' not found in {root}")
dicoms_dir = inbreast_dir / "AllDICOMs"
xml_dir = inbreast_dir / "AllXML"
roi_dir = inbreast_dir / "AllROI"
if not dicoms_dir.exists():
raise FileNotFoundError(f"AllDICOMs not found at {dicoms_dir}")
# Parse CSV metadata (reuse INbreast logic)
self._csv_metadata = INbreastDataset._parse_csv(inbreast_dir)
# Pair DICOMs ↔ XMLs ↔ Masks
self._samples: List[CustomINbreastSample] = pair_custom_inbreast(
str(dicoms_dir),
str(xml_dir),
str(masks_dir) if masks_dir.exists() else "",
str(roi_dir) if roi_dir.exists() else None,
)
logger.info(
f"CustomINbreastDataset: {len(self._samples)} samples, "
f"{sum(1 for s in self._samples if s.mask_path)} with TIF masks"
)
# ------------------------------------------------------------------
# Load single sample
# ------------------------------------------------------------------
def _load_sample(self, idx: int) -> Dict[str, Any]:
sample: CustomINbreastSample = self._samples[idx]
# Load DICOM
dcm = DicomImage(file_path=sample.dicom_path)
dcm.load()
image_tensor = self._to_chw(dcm.pixel_data.float())
h, w = dcm.height, dcm.width
mask_source = "empty"
# Priority 1: TIF mask
if sample.mask_path is not None:
try:
mask_np = load_tif_mask(sample.mask_path)
# Resize mask if dimensions don't match DICOM
if mask_np.shape != (h, w):
logger.debug(
f"Resizing mask {mask_np.shape} → ({h}, {w}) "
f"for case {sample.case_id}"
)
mask_tensor = (
torch.from_numpy(mask_np).unsqueeze(0).unsqueeze(0).float()
)
mask_tensor = torch.nn.functional.interpolate(
mask_tensor, size=(h, w), mode="nearest"
).squeeze(0)
else:
mask_tensor = torch.from_numpy(mask_np).unsqueeze(0).float()
mask_source = "tif"
except Exception as e:
logger.warning(
f"Failed to load TIF mask for case {sample.case_id}: {e}"
)
mask_tensor = None
# Priority 2: XML fallback
if mask_source != "tif" and sample.xml_path is not None:
try:
mask_np = xml_to_binary_mask(
sample.xml_path, (h, w), point_radius=self.point_radius
)
mask_tensor = torch.from_numpy(mask_np).unsqueeze(0).float()
mask_source = "xml"
except Exception as e:
logger.warning(
f"Failed to generate XML mask for case {sample.case_id}: {e}"
)
mask_tensor = None
# Fallback: empty mask
if mask_source == "empty" or mask_tensor is None:
mask_tensor = torch.zeros(1, h, w, dtype=torch.float32)
mask_source = "empty"
logger.warning(f"No mask available for case {sample.case_id}")
# Build metadata
csv_meta = self._csv_metadata.get(sample.case_id, {})
metadata = {
"case_id": sample.case_id,
"mask_source": mask_source,
"laterality": csv_meta.get("laterality", ""),
"view": csv_meta.get("view", ""),
"birads": csv_meta.get("birads", -1),
"file_name": Path(sample.dicom_path).name,
}
return {
"image": image_tensor,
"mask": mask_tensor,
"metadata": metadata,
}