Datasets

PyTorch-compatible dataset classes for medical imaging research.

BaseDataset

class medical_image.datasets.base_dataset.BaseDataset[source]

Bases: Dataset, ABC

Abstract base class for PyTorch-compatible medical image datasets.

Enforces lazy loading: images are never pre-loaded into memory. Each sample is loaded on-the-fly in __getitem__.

Subclasses must implement:

  • _build_sample_list(): scan the dataset directory and build an internal list of samples (paths, metadata).

  • _load_sample(idx): load a single sample by index, returning a dict with "image", "mask" (or "label"), and "metadata".

Parameters:

  • root_dir – Root directory of the dataset.

  • transform – Optional callable transform applied to the image tensor.

  • target_transform – Optional callable transform applied to the mask/label.

  • target_size – Optional (H, W) to resize images and masks on load.

__init__(root_dir, transform=None, target_transform=None, target_size=None)[source]
Parameters:
classmethod download(source, destination, method='local', percentage=None)[source]

Download the dataset from a source.

Parameters:

  • source (str) – Source path or URL.

  • destination (str) – Local destination directory.

  • method (Literal['local', 'http', 'ftp']) – 'local', 'http', or 'ftp'.

  • percentage (float | None) – Optional subset percentage (0–1] for large datasets.

Returns:

Absolute path to the downloaded dataset.

Return type:

str

to_coco_json(output_path=None, description='Medical Image Dataset')[source]

Export the entire dataset as a COCO-format JSON.

Iterates over every sample, calls _get_annotations() for each, and builds the standard COCO structure (info, images, annotations, categories).

Each annotation entry includes a custom "center" field with the annotation centroid – this is an extension to the official COCO spec.

Note

COCO bbox format is [x, y, width, height], not [x_min, y_min, x_max, y_max]. The conversion is handled automatically.

Parameters:

  • output_path (str | None) – If provided, the JSON dict is also written to this file path.

  • description (str) – Free-text description for the COCO info block.

Returns:

The full COCO JSON structure as a dict.

Return type:

dict

classmethod from_coco_json(json_path)[source]

Load dataset metadata and annotations from a COCO JSON file.

Reconstructs Annotation objects from COCO segmentation polygons (preferred) or bounding boxes (fallback when segmentation is empty).

Parameters:

json_path (str) – Path to a COCO-format .json file.

Returns:

  • "images" – list of COCO image-entry dicts.

  • "annotations" – Dict[int, List[Annotation]] mapping each COCO image ID to its reconstructed annotations.

  • "categories" – Dict[int, str] mapping each COCO category ID to its label name.

Return type:

A dict with three keys

INbreastDataset

class medical_image.datasets.inbreast.INbreastDataset[source]

Bases: BaseDataset

PyTorch Dataset for the INbreast mammography database.

Supports two directory layouts:

Layout A — COCO annotations (preferred):

root_dir/
├── annotations.json     # COCO format
└── images/
    ├── 20586934.dcm
    └── ...

Layout B — XML annotations (legacy):

root_dir/
├── AllDICOMs/
├── AllXML/
├── AllROI/       (optional)
└── INbreast.csv  (optional)

Each sample returns:

{
    "image": Tensor[1, H, W],
    "mask":  Tensor[1, H, W],
    "metadata": {
        "case_id": str,
        "file_name": str,
        "num_annotations": int,
    }
}
Parameters:

  • root_dir – Path to the dataset root directory.

  • transform – Optional transform for the image tensor.

  • target_transform – Optional transform for the mask tensor.

  • target_size – Optional (H, W) resize target.

  • point_radius – Radius for rendering single-point ROIs (XML mode only).

Example:

dataset = INbreastDataset("data/Inbreast")
sample = dataset[0]
print(sample["image"].shape)     # torch.Size([1, 512, 512])
print(sample["mask"].shape)      # torch.Size([1, 512, 512])
__init__(root_dir, transform=None, target_transform=None, target_size=None, point_radius=3)[source]
Parameters:

CustomINbreastDataset

class medical_image.datasets.custom_inbreast.CustomINbreastDataset[source]

Bases: 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:

{
    "image": Tensor[1, H, W],
    "mask":  Tensor[1, H, W],     # from TIF or XML fallback
    "metadata": {"case_id": str, "mask_source": "tif"|"xml"|"empty", ...}
}
Parameters:

  • 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.

__init__(root_dir, transform=None, target_transform=None, target_size=None, point_radius=3)[source]
Parameters:

CBISDDSMDataset

class medical_image.datasets.cbis_ddsm.CBISDDSMDataset[source]

Bases: BaseDataset

PyTorch Dataset for the CBIS-DDSM mammography database.

Supports two loading modes:

  • "full_image": Load the entire mammogram (optionally resized).

  • "patch": Extract sliding-window patches from each mammogram.

Each sample via __getitem__ returns:

{
    "image": Tensor[1, H, W],      # mammogram (or patch)
    "mask":  Tensor[1, H, W],       # OR-merged ROI masks
    "metadata": {
        "case_id": str,
        "patient_id": str,
        "side": str,
        "view": str,
        "num_masks": int,
        "patch_idx": int,           # only in patch mode
        "patch_position": (y, x),   # only in patch mode
    }
}

Use get_detailed_sample() for the rich format with per-ROI bounding boxes, crops, and masks.

Parameters:

  • root_dir – Path to the manifest directory containing CBIS-DDSM/.

  • mode – "full_image" or "patch".

  • patch_size – Patch side length (used when mode="patch").

  • stride – Stride between patches (used when mode="patch").

  • transform – Optional image transform.

  • target_transform – Optional mask transform.

  • target_size – Optional (H, W) resize for full_image mode.

  • percentage – Optional float (0-1] to use a random subset of cases.

  • seed – Random seed for reproducible subset selection.

Example:

dataset = CBISDDSMDataset(
    root_dir="data/ddsm",
    mode="patch",
    patch_size=512,
    stride=256,
    percentage=0.2,
)

loader = DataLoader(
    dataset,
    batch_size=4,
    shuffle=True,
    num_workers=4,
    collate_fn=dataset.collate_fn,
)
__init__(root_dir, mode='full_image', patch_size=512, stride=256, transform=None, target_transform=None, target_size=None, percentage=None, seed=42)[source]
Parameters:
get_detailed_sample(idx)[source]

Load a full-image sample with per-ROI bounding boxes, crops, and masks.

Only available in "full_image" mode.

Returns:

  • "image": Tensor[1, H, W] — full mammogram

  • "bboxes": Tensor[N, 4] — (x_min, y_min, x_max, y_max) in full mammogram coordinates

  • "rois": List[Tensor] — ROI crop tensors (from 1-1.dcm)

  • "masks": List[Tensor] — ROI mask tensors (from 1-2.dcm)

  • "label": task string (e.g. "Calc-Test")

  • "meta": dict with patient_id, view, side, task

Return type:

Dict with

Parameters:

idx (int)

static get_bounding_boxes(mask)[source]

Compute bounding boxes from a binary mask array.

Finds connected components of non-zero pixels and returns the bounding box of each component.

Parameters:

mask (ndarray) – 2D numpy array (H, W) with non-zero pixels marking ROIs.

Returns:

List of (x_min, y_min, x_max, y_max) tuples.

Return type:

List[Tuple[int, int, int, int]]

static collate_fn(batch)[source]

Custom collate function that handles variable-sized masks.

Pads all images and masks to the maximum spatial dimensions in the batch.

Returns:

Dict with stacked "image" and "mask" tensors, and a list of "metadata" dicts.

Parameters:

batch (List[Dict[str, Any]])

Return type:

Dict[str, Any]