From DICOM to Diagnosis: Five Engineering Problems Behind a Mammography AI Pipeline
Introduction
When hospitals provide radiologists with an AI tool, it relies on several software steps built around the neural network:
the image needs to be read as the scanner intended
the labels used for training must match the lesions they describe.
any mark shown to the viewer should be placed over the tissue the model actually analyzed, and it should stay connected to its score.
If one of the above fails, the software can still produce results that look reasonable, so the error might not be noticed.
This document explains how we build and test that chain. Our example is a research pipeline for breast lesion analysis that we developed from the ground up. It reads mammograms in DICOM format, finds suspicious areas, outlines masses, estimates the chance that each lesion is malignant, and returns results as standard DICOM objects that can be viewed with the original image.
The pipeline is trained on CBIS-DDSM, a public set of about 10,000 scanned-film mammograms with lesion outlines by radiologists and biopsy-confirmed pathology. Anyone who downloads the data can reproduce every step.
We chose three well-known networks: a Faster R-CNN detector, a U-Net for segmentation, and a ResNet50 classifier, each starting from published pretrained weights. By using familiar models, we could focus on the engineering around them, which is the main topic of this article. The five sections follow an image as it moves through the pipeline, covering data integrity, image correctness, geometry preservation, interoperability, and verification. Each section describes possible problems, what we built to prevent them, the code involved, and test results showing if the solution works. Some sections also share mistakes we made and how we found them, since this gives readers a better sense of how a team operates than just showing a finished design.
The pipeline is fully implemented and has been tested on synthetic DICOM data designed to match the structure and quirks of CBIS-DDSM. Using synthetic data with known answers lets us check if a mask matches the correct lesion or if a box covers the right pixels, and we report all those results. To measure how well the models detect and classify real cancers, we need real images. The models have not yet been trained or tested on real data, so this document does not measure diagnostic accuracy.
The pipeline at a glance
A detector scans the whole breast to find possible lesions. Each candidate is cropped from the full-resolution image; then a classifier estimates the likelihood of malignancy, and a U-Net outlines the masses. The results are overlaid on the original image and saved as DICOM files. The networks used are a Faster R-CNN with a ResNet50 backbone, a U-Net with a ResNet34 encoder, and a ResNet50 classifier, all starting from ImageNet weights. These are standard, well-known choices, so this article does not focus on them.
1. Data integrity: proving a label belongs to its lesion
In CBIS-DDSM, each lesion is described in a CSV file that references three DICOM files: the full mammogram, a binary mask showing the lesion's location, and a cropped image of the lesion. However, the file paths in the CSV often differ from the actual downloaded folder structure. To avoid this problem, the pipeline identifies files using their DICOM PatientID instead of relying on file paths. This ID contains the case, breast, and view information, as well as the abnormality number for lesion-specific files.
This approach helps locate the files, but it does not guarantee they are the right ones. For example, a series labeled as a mask may also contain a cropped image. Some masks may be incomplete, empty, or associated with the wrong lesion. The most serious problem occurs when a mask looks valid but belongs to a different lesion, causing the model to learn from incorrect data. To prevent this, the pipeline accepts a mask only when it passes five independent checks:
Evidence | Requirement |
|---|---|
Identifier | PatientID carries the CSV row's abnormality number |
Metadata | Breast side, view and lesion type in the key agree with the CSV row |
Dimensions | The mask matches the full image's size, within 16 pixels |
Content | Binary, non-empty, under half the image |
Pixels | The lesion crop from the same series is found in the full image, around the mask |
The first four checks identify invalid or incomplete files. However, only the pixel check can detect a valid mask that belongs to the wrong lesion. We improved this check twice to make it more reliable.
The initial version compared raw pixel values using normalized cross-correlation, but it treated strong negative and positive correlations equally. This meant an incorrect region could appear to be a good match.
The current version applies the same DICOM grayscale processing to both the cropped lesion and the full mammogram, using each file's metadata to interpret pixel values correctly. It then looks for a strong positive match that is clearly better than matches elsewhere in the image. Finally, it verifies that the mask lies within the matched crop and that the lesion is positioned near the crop's center.
To test the matcher, we introduced ten types of faults into 1,778 synthetic abnormalities. Three of these faults were specifically designed to defeat the pixel check. We then compared different configurations:
Case | No pixel check | First pixel check | Current |
|---|---|---|---|
Correct pairs accepted | 1,156 of 1,156 | 1,156 of 1,156 | 1,156 of 1,156 |
Masks swapped between two lesions | 0 of 88 rejected | 87 of 88 | 88 of 88 |
Mask borrowed from another patient | 0 of 62 rejected | 61 of 62 | 61 of 62 |
Inverted lesion copy planted elsewhere, mask moved onto it | 0 of 23 rejected | 0 of 23 | 23 of 23 |
Crop inverted while its header says otherwise | 0 of 69 rejected | 0 of 69 | 69 of 69 |
Crop and image spacing differ | 0 of 26 rejected | 0 of 26 | 26 of 26 |
Malformed masks and CSV mismatches | 354 of 354 rejected | 354 of 354 | 354 of 354 |
One incorrect mask still passed the checks. It belonged to another patient whose lesion was located in almost the same position. This revealed a limitation: the pixel check can confirm that a lesion exists in the expected area, but cannot guarantee that the mask outlines the correct lesion.
Synthetic crops are always perfectly aligned with their lesions, while real CBIS-DDSM crops may be misaligned. Therefore, you may need to adjust matching thresholds when working with real data. To support this, the pipeline records matching scores for accepted masks and reasons for rejecting others, allowing thresholds to be adjusted based on actual results.
The train/test split also requires attention. CBIS-DDSM originally separated training and test data for masses and calcifications independently, meaning images from the same patient can appear in both sets. To prevent this data leakage, the pipeline checks patient IDs across all four subsets. By default, if a patient appears in any official test subset, we assign all their images to the test set. The pipeline then selects validation patients from the remaining training patients. Every metrics file records the splitting protocol used, making the evaluation results traceable.
2. Image correctness: what the model actually sees
A DICOM file has both pixel data and metadata that tell us how to interpret the pixels. Our pipeline handles this in three steps. First, it decodes the pixel data, no matter how it was stored. Next, it applies DICOM grayscale transformations, like intensity adjustments, optional windowing, and image polarity. Finally, it normalizes the pixel values to the [0, 1] range for the ML model. This pipeline does not include other operations for clinical display, such as calibration, shutters, and overlays.
Image polarity turned out to be a key source of error. In our setup, higher pixel values mean denser tissue, which looks brighter in the image. DICOM uses the Photometric Interpretation attribute to describe polarity. For example, MONOCHROME1 means the lowest pixel value appears white.
In our first version, we treated Presentation LUT Shape = INVERSE as a signal to invert the image again. This was a mistake. The DICOM DX Image Module specification, used for mammography images, says MONOCHROME1 and INVERSE describe the same polarity, not two separate inversions.
As a result, we processed valid MONOCHROME1 mammograms with reversed intensity, so dense tissue looked dark instead of bright. The incorrectly processed images had a correlation of −1.000 with the correct output, showing a complete reversal of pixel intensity.
Surprisingly, our automated test passed. We had written the test using the same wrong understanding of the DICOM standard, so it confirmed the mistake instead of catching it. A reviewer finally found the problem by comparing our implementation to the standard.
This taught us an important lesson: automated tests can check that code does what you expect, but they cannot make sure those expectations are right. Independent review against the specification is just as important.
Now, we get polarity only from Photometric Interpretation, and use Presentation LUT Shape as a consistency check:
Photometric Interpretation | Presentation LUT Shape | Model input |
|---|---|---|
MONOCHROME2 | Absent or IDENTITY | Stored order kept |
MONOCHROME1 | Absent or INVERSE | Inverted |
MONOCHROME2 | INVERSE | Rejected: contradicts the module definition |
MONOCHROME1 | IDENTITY | Rejected: contradicts the module definition |
The Modality LUT is applied before the VOI transformation, following pydicom's processing order, because windowing operates on modality-adjusted pixel values. VOI processing is disabled by default, since vendor-defined window settings are intended for radiologists, while the ML model benefits from consistent pixel values. The same processing function also matches lesion crops with full mammograms, so the polarity correction fixed that part of the pipeline as well.
We also expanded the encoding test. It now creates 19 versions of the same synthetic mammogram, covering seven transfer syntaxes, combinations of polarity settings, rescaling, Modality LUT, signed pixel values, 8-bit storage, and VOI windowing. All lossless versions produced identical model inputs.
The pipeline also correctly rejected five unsupported or invalid cases: color images, multi-frame images, lossy compression, and two conflicting polarity combinations.
3. Geometry preservation: returning to the original pixels
The detector does not work directly with the original DICOM image. Each mammogram is first flipped, if needed, so the chest wall is on the left. Labels and other unwanted areas are removed, the image is cropped to the breast, and then resized for the model. Lesion patches are extracted separately from the flipped image at its original resolution, preserving tiny details such as microcalcifications that may be only a few pixels wide.
Because the pipeline uses images at different sizes and orientations, it must map predictions back to their exact locations in the original DICOM. To do this, it records whether the image was flipped, where it was cropped, and how much it was resized. We can then reverse these transformations for every predicted bounding box or mask.
We verified this process using synthetic data. All 158 lesion boxes, including 54 from flipped images, were transformed into model coordinates and back to their original positions with zero pixel error.
The more challenging part is correctly identifying the breast region. We use Otsu thresholding and connected-component analysis to separate the breast from the background. Additional processing removes labels that touch the breast boundary and includes faint skin edges that might otherwise be missed.
This approach works well for clean scanned-film images, but may not perform equally well on mammograms from different clinical vendors. For this reason, we check every generated breast mask. If the check fails, the pipeline keeps the full image without masking or cropping, rather than risking the removal of important tissue.
Preprocessing adds one more check with the training labels. If any lesion lies partly outside the breast mask, masking would erase part of it, so that image falls back too.
We stress-tested the heuristic with 200 synthetic images for each of seven conditions, each reproducing a known failure mode, and counted failures the checks did not flag:
Condition | Failures, first version (of 60) | Failures, current (of 200) | Silent failures, current |
|---|---|---|---|
Clean film | not measured | 0 | 0 |
Low-contrast skin edge | 60 | 0 | 0 |
Lesion near a low-contrast skin edge | 60 | 0 | 0 |
Label touching the breast | 53 | 0 | 0 |
Bright pectoral muscle | not measured | 0 | 0 |
Scanner line reaching the far edge | not measured | 0 | 0 |
Film fog brightening the background | not measured | 1 | 0 |
The pipeline detected the one failure caused by image fog, but the fallback method also incorrectly estimated the breast orientation. As a result, the image was kept at its original size but remained mirrored. We recorded this issue in the quality report for review.
Fog also caused 88 unnecessary fallbacks. In these cases, the images were left uncropped, but no image content was lost.
The actual failure rate on CBIS-DDSM is still unknown. It will be measured once the full dataset has been processed, using the same quality report to identify and assess any problems.
4. Interoperability: keeping findings identifiable
Radiologists rely on their own DICOM viewers, so the pipeline needs to create results that these systems can show without extra steps. It produces three outputs: a DICOM Segmentation (SEG) with lesion outlines, a Secondary Capture image with annotated findings, and a JSON file for integration with other systems.
The SEG links back to the original mammogram using its SOP Instance UID. This lets compatible viewers like OHIF, 3D Slicer, and multiple PACS systems show the lesion outlines on the source image.
The main challenge is keeping each lesion identifiable across all three outputs. In case multiple detected masses are combined into a single SEG segment, it means two lesions in the same image appeared as one combined region, making it impossible to link each lesion to its individual confidence score.
The correct pipeline should create a separate segment for every lesion, even when their outlines overlap. Each lesion receives a unique Tracking ID that corresponds to its finding number in the JSON file and annotated image. This lets you trace every finding and its associated results across all three outputs.
In the DICOM SEG, each lesion has its own layer in the pixel array and a corresponding description. These must appear in the same order. Otherwise, a lesion outline could be assigned the wrong label without any warning.
Calcification clusters are handled differently. They appear as detections in the JSON file and annotated image, with a note explaining that they have not been segmented. We plan to add support for DICOM Structured Reports (TID 1500), which can describe each calcification cluster using a bounding region.
De-identified research datasets introduce another challenge. CBIS-DDSM files may be missing patient or study information, such as birth dates. DICOM requires certain fields, known as Type 2 attributes, to be present even if their values are empty. Since highdicom cannot create a SEG without these fields, the pipeline adds them with empty values to a copy of the original DICOM header.
We use two tests to verify the SEG output. The first creates several known lesion masks, including two that overlap, and then reads the generated SEG back. Every segment must match its original mask pixel for pixel, with the correct label, property type, and Tracking ID.
The second test runs the full inference pipeline and compares the SEG output with the JSON results. It verifies that each segment corresponds to exactly one detected mass, has the same number of pixels, and that no calcification detections are incorrectly included in the SEG.
5. Verification: what has been tested, and what has not
Claims about medical AI software need two types of evidence, each answering a different question.
Engineering verification checks whether the software works as intended. This includes reading DICOM files correctly, matching each lesion to the right label, and ensuring detected findings appear in the correct locations in the original images and output files.
Model evaluation measures how accurately the trained models detect and classify cancers in real medical images.
Engineering verification can be performed using synthetic data, where we know exactly what the correct results should be. Model evaluation, however, requires real medical images because synthetic lesions are generally much easier to detect than real ones.
This section presents the engineering verification results in detail and explains what still needs to be done to evaluate the models on real clinical data.
Engineering verification: completed
The table below summarizes the tests used to verify the fixes described in Sections 1-4. All tests used synthetic DICOM files designed to reproduce the structure and known issues of CBIS-DDSM. The tests can be rerun using the code in the repository.
Passing these tests confirms that the software works as intended under the tested conditions. It does not demonstrate how accurately the models detect or classify cancer in real medical images.
Area | What was tested | Result |
|---|---|---|
ROI matching | 1,778 abnormalities, 622 with injected faults | All correct pairs accepted; 621 of 622 faults rejected |
Patient split | Mass and calcification splits drawn separately | Audit flagged the 2 of 30 patients spanning train and test; none overlap under the default protocol |
Pixel decoding and polarity | One image in 19 variants, plus 5 files to refuse | Identical input for every lossless variant; all 5 refused |
Breast region | 1,400 stress images in 7 conditions | 1 failure, flagged; no silent failures |
Coordinates | 158 lesion boxes, original to model and back | 0 pixel error |
DICOM SEG | Known masks, and inference output against its JSON | One segment per lesion, pixel-identical, correct identities |
Metric code | FROC, AP and Dice on hand-built cases | Matched hand-computed values |
Article code | Verbatim excerpts compared with the repository | All match, checked automatically |
Model evaluation: still to come
The models have not yet been trained or tested on real CBIS-DDSM images, so we cannot report diagnostic accuracy at this stage.
Evaluation is planned in two stages. First, we will test the models on CBIS-DDSM using a patient-level split, ensuring that images from the same patient never appear in both training and test sets. This avoids the data overlap issue described in Section 1.
Second, we will evaluate the models on an independent digital mammography dataset, such as CMMD, VinDr-Mammo, or INbreast. This is important because CBIS-DDSM contains scanned-film mammograms, while modern clinical systems use digital mammography. Testing on a separate dataset will help assess how well the models perform on images from different sources.
Each task will be evaluated using metrics appropriate to what the model is expected to do:
Task | Metric |
|---|---|
Detection | FROC sensitivity at 0.25 to 4 false positives per image |
Mass segmentation | Dice |
Malignancy classification | AUROC, and sensitivity and specificity at a threshold chosen on validation |
How the tests were designed
Reliable testing depends on knowing the correct result. We followed three principles to make our tests trustworthy.
First, we derive expected results from the DICOM standard, not our own code. For example, the polarity tests use the DX Image Module specification to determine how to interpret pixel values. This helped us identify the mistake described in Section 2, where both the implementation and its original test followed the same incorrect interpretation.
Second, we create test files independently of the pipeline. We encode the same image using seven DICOM transfer syntaxes, with pydicom handling most encodings and GDCM handling JPEG Lossless. This lets us verify that the pipeline correctly reads files it did not generate. We also read every generated DICOM SEG back and compare its contents pixel for pixel with the expected lesion masks.
Third, we report actual test counts rather than relying only on percentages. For example, reporting 621 successful cases out of 622 makes the single failure visible. A percentage rounded to 0% could hide an error that matters in clinical applications.
What this means for clinical AI software
The five problems discussed here have one thing in common: the software gave results that looked correct but were not. None of these issues caused a crash, and some might have been missed even when checking model evaluation metrics.
Our experience with this pipeline showed that the most effective way to catch such errors is to identify assumptions, turn them into automated checks, and deliberately test those checks with cases designed to expose weaknesses.
Several steps remain before this pipeline could become a clinical product. First, the models must be trained and evaluated on real CBIS-DDSM images, then tested on an independent digital mammography dataset. This is important because CBIS-DDSM contains scanned-film images, while modern mammography systems use digital detectors.
A clinical system would also need to analyze both views of both breasts, ideally including previous examinations, rather than treating each image independently.
To fit into clinical workflows, the system could use a DICOM server like Orthanc or a deployment framework like MONAI Deploy, following the IHE AI Results profile to share findings with existing imaging systems.
Finally, AI software used to support medical diagnosis is generally regulated as a medical device. Moving toward clinical use would require an appropriate quality management system, compliance with relevant standards such as ISO 13485 and IEC 62304, and applicable regulatory authorization.
The main takeaway is that it is not enough for a model to make predictions that seem reasonable. Clinical AI software also needs to show that its data processing, image geometry, outputs, and evaluation methods are correct, reliable, and can be checked independently.
Appendix: Model details
This appendix brings together the model-specific details from the main article. It explains how we measure each network’s quality and how we build and train the detector, segmentation network, and classifier. The code shown here is shortened for easier reading.
Evaluation metrics
Report the metric each task actually needs.
Task | Metric | Why this one |
|---|---|---|
Detection | FROC: sensitivity at 0.25, 0.5, 1, 2 and 4 false positives per image | Radiologists tolerate a fixed number of false marks per image; accuracy is meaningless when most pixels are background |
Segmentation | Dice coefficient | Measures overlap of the predicted outline with the radiologist's |
Classification | AUROC, plus sensitivity and specificity at a fixed threshold | Shows the trade-off between missed cancers and unnecessary biopsies |
The classification threshold is chosen on validation, using Youden's J, then frozen and applied unchanged to the test set. Tuning it on test data is a common way results end up looking better than they will in practice.
Finding lesions: the detector
The detector helps us find suspicious regions across the entire mammogram. We use torchvision's Faster R-CNN with a ResNet50 backbone and a Feature Pyramid Network (FPN). The FPN is important because it predicts from several feature-map resolutions, making it possible to detect both a 2 cm mass and a 3 mm calcification cluster.
Converting masks to detection targets
CBIS-DDSM gives masks instead of boxes. To get a box, use the smallest rectangle that covers each mask and map it to the preprocessed image's coordinates using the geometry record.
def mask_box(mask):
ys, xs = np.where(mask)
# x1, y1, x2, y2
return [xs.min(), ys.min(), xs.max() + 1, ys.max() + 1]
# flip like the image
m = canonical_mask(read_mask(row["mask_dcm"], img.shape), geo)
# crop offset, then scale
box = geo.canonical_to_model(mask_box(m))
Assign a class label to each box: use 1 for mass and 2 for calcification cluster. In torchvision, 0 is reserved for background.
The dataset class
Torchvision detection models use a list of 3-channel images and a list of target dictionaries. Grayscale images are converted to three channels so that ImageNet-pretrained weights can be used.
class MammoDetDataset(Dataset):
def __getitem__(self, i):
r = self.recs[i]
# 16-bit PNG -> [0, 1]
img = read_png01(self.dir / r["image"])
boxes = np.array([l["box"] for l in r["lesions"]], np.float32)
labels = np.array([LESION_TYPES[l["type"]] for l in r["lesions"]], np.int64)
if self.train:
img = np.clip(img ** random.uniform(0.7, 1.4) * random.uniform(0.9, 1.1), 0, 1)
# vertical flip only
if random.random() < 0.5:
img = img[::-1].copy()
boxes[:, [1, 3]] = img.shape[0] - boxes[:, [3, 1]]
x = torch.from_numpy(np.repeat(img[None], 3, axis=0))
return x, {"boxes": torch.from_numpy(boxes), "labels": torch.from_numpy(labels)}
We used only mild augmentation. Adjusting gamma and contrast simulates different scanners and exposures, and vertical flips are safe for anatomy. We did not use horizontal flips because preprocessing already standardizes the chest wall side.
The model
from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
def build_detector(num_classes=3, long_side=1536):
model = fasterrcnn_resnet50_fpn_v2(
weights="DEFAULT", # COCO-pretrained
min_size=long_side, max_size=long_side,
box_detections_per_img=50)
in_f = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_f, num_classes)
return modelTwo implementation details are important here. First, replacing box_predictor changes the model's original COCO classification head from 91 classes to our 3 classes, while preserving the pretrained backbone and region proposal network.
Second, setting min_size and max_size to match the preprocessed image size prevents torchvision from resizing the image again. This is particularly important for detecting small calcifications, which could become even harder to detect if the image were reduced further.
Training
During training, a torchvision detector returns several loss values: how well it identifies potential object regions, how accurately it adjusts their bounding boxes, how well it classifies detections, and how accurately it refines the final boxes. These losses are added together to calculate the total training loss.
opt = torch.optim.SGD(params, lr=0.01, momentum=0.9, weight_decay=1e-4)
warmup = torch.optim.lr_scheduler.LinearLR(opt, 0.1, total_iters=500)
cosine = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
scaler = torch.amp.GradScaler("cuda")
for epoch in range(epochs):
model.train()
for xs, ts in loader:
xs = [x.to(device) for x in xs]
ts = [{k: v.to(device) for k, v in t.items()} for t in ts]
with torch.autocast("cuda"):
loss = sum(model(xs, ts).values())
opt.zero_grad()
scaler.scale(loss).backward()
scaler.unscale_(opt)
# stabilizes early epochs
torch.nn.utils.clip_grad_norm_(params, 10.0)
scaler.step(opt); scaler.update()
if epoch == 0: warmup.step()
cosine.step()The training loop uses SGD optimization, an initial learning-rate warmup, and cosine learning-rate scheduling. Gradient clipping helps prevent excessively large updates, especially during early training.
Mixed precision training (autocast and GradScaler) reduces GPU memory usage, which is particularly useful when processing mammograms at 1,536-pixel resolution.
The default batch size is 4 images, although the optimal value depends on the available GPU memory and hardware. Training performance and memory requirements have not yet been measured on real CBIS-DDSM images.
Evaluating with FROC
Each prediction is matched greedily, highest score first, to an unmatched ground-truth lesion with IoU of at least 0.3. Then sensitivity is read off at fixed false-positive rates per image:
def froc(scores, is_tp, n_lesions, n_images, fp_rates=(0.25, 0.5, 1, 2, 4)):
order = np.argsort(-scores)
tp, fp = np.cumsum(is_tp[order]), np.cumsum(~is_tp[order])
sens, fppi = tp / n_lesions, fp / n_images
return {f"sens@{r}": sens[fppi <= r].max() if (fppi <= r).any() else 0.0
for r in fp_rates}The checkpoint with the best mean FROC sensitivity on validation is kept. An IoU of 0.3 rather than COCO's 0.5 is common in mammography, because the boundary of a mass or cluster is itself uncertain.
One property of CBIS-DDSM limits what these numbers mean. Every image contains at least one lesion, so the detector never sees a normal mammogram, and its measured false-positive rate is optimistic. A realistic estimate needs normal screening images from another source in the validation set.
Outlining masses: the segmentation network
Once a mass is found, the segmentation network traces its exact outline. It works on a patch centred on each lesion, which keeps the fine boundary detail at full resolution.
Cutting lesion patches
Each patch is a square centred on the lesion, 1.5 times its size with a minimum of 512 original pixels, cut from the full-resolution canonical image and resized to 256 × 256. The same function cuts the matching mask:
def patch_window(box, context=1.5, min_side=512):
x1, y1, x2, y2 = box
side = int(max(max(x2 - x1, y2 - y1) * context, min_side))
cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
return int(cx - side / 2), int(cy - side / 2), side
def extract_patch(img, window, out_size=256):
x0, y0, side = window
# zero-pad at image borders
patch = np.zeros((side, side), img.dtype)
sx0, sy0 = max(x0, 0), max(y0, 0)
sx1, sy1 = min(x0 + side, img.shape[1]), min(y0 + side, img.shape[0])
patch[sy0 - y0:sy1 - y0, sx0 - x0:sx1 - x0] = img[sy0:sy1, sx0:sx1]
return cv2.resize(patch, (out_size, out_size), interpolation=cv2.INTER_AREA)
The context factor gives the network surrounding tissue to compare against, which is how radiologists judge margins too. The minimum size stops tiny lesions from being blown up into blurry patches.
The architecture: U-Net with a pretrained encoder
U-Net is a neural network designed for image segmentation. It has two main parts: an encoder, which extracts important features while reducing the image resolution, and a decoder, which uses those features to create a pixel-level mask showing the lesion's shape.
Skip connections link the encoder and decoder, helping preserve fine details that might otherwise be lost during processing.
We use a pretrained ResNet34 as the encoder. Because it has already learned useful image features, the model can generally learn faster than a U-Net trained entirely from scratch:
class Up(nn.Module):
def __init__(self, cin, cskip, cout):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(cin + cskip, cout, 3, padding=1, bias=False), nn.BatchNorm2d(cout), nn.ReLU(True),
nn.Conv2d(cout, cout, 3, padding=1, bias=False), nn.BatchNorm2d(cout), nn.ReLU(True))
def forward(self, x, skip=None):
x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=False)
return self.conv(torch.cat([x, skip], 1) if skip is not None else x)
class UNetResNet34(nn.Module):
def __init__(self):
super().__init__()
enc = torchvision.models.resnet34(weights="DEFAULT")
# 1/2, 64 ch
self.stem = nn.Sequential(enc.conv1, enc.bn1, enc.relu)
# 1/4
self.pool = enc.maxpool
self.l1, self.l2, self.l3, self.l4 = enc.layer1, enc.layer2, enc.layer3, enc.layer4
self.d4, self.d3 = Up(512, 256, 256), Up(256, 128, 128)
self.d2, self.d1 = Up(128, 64, 64), Up(64, 64, 64)
self.d0, self.head = Up(64, 0, 32), nn.Conv2d(32, 1, 1)
def forward(self, x):
s0 = self.stem(x); s1 = self.l1(self.pool(s0))
# down to 1/32
s2 = self.l2(s1); s3 = self.l3(s2); s4 = self.l4(s3)
x = self.d4(s4, s3); x = self.d3(x, s2); x = self.d2(x, s1); x = self.d1(x, s0)
return self.head(self.d0(x))
The loss
A lesion usually occupies only a small area of an image patch, while most pixels belong to the background. This creates a problem during training: using cross-entropy loss alone may encourage the model to predict background for nearly every pixel.
To address this, we combine cross-entropy with soft Dice loss, which measures how closely the predicted lesion mask overlaps with the correct mask. This helps the model focus on identifying the lesion rather than simply predicting the background:
def dice_bce(logits, target):
p = torch.sigmoid(logits)
inter = (p * target).sum((1, 2, 3))
dice = 1 - (2 * inter + 1) / (p.sum((1, 2, 3)) + target.sum((1, 2, 3)) + 1)
return F.binary_cross_entropy_with_logits(logits, target) + dice.mean()
Training and inference
Because lesion patches have no fixed orientation, we can rotate and flip them in eight different ways during training. We also apply small changes in scale and rotation to help the model recognize lesions under slightly different conditions.
Training uses the AdamW optimizer with a learning rate of 3e-4, cosine learning-rate decay, and a batch size of 16. We save the model checkpoint that achieves the highest Dice score on the validation set.
During inference, the lesion patch is extracted using the bounding box predicted by the detector, rather than the known lesion location. The segmentation model produces a mask, which is converted to a binary image using a threshold of 0.5. The mask is then mapped back to its correct position in the original full-resolution mammogram by reversing the patch extraction transformation.
Why do we segment masses only? In CBIS-DDSM, level-set segmentation was used for masses, while calcification masks represent entire clusters rather than individual tiny calcifications. These cluster masks are therefore less suitable for training a model to identify precise pixel-level boundaries.
One limitation remains. During training, patches are centered on the actual lesion. During inference, however, they are centered on the detector's predicted bounding box, which may be slightly inaccurate.
To make the segmentation model more tolerant of these errors, we plan to randomly shift patch centers during training by approximately 10% of the patch size. This would help the model handle imperfect detections and is our first planned improvement.
Judging malignancy: the classifier
The final model estimates the probability that a detected lesion is malignant (cancerous). It uses the same lesion patches as the segmentation model, allowing both models to share the same preprocessing steps.
Why patches instead of the whole image?
A malignant mass may occupy as little as 1% of a mammogram. Reducing a full 3,000 × 5,000-pixel image to the size required by a neural network can remove important details, such as lesion edges and tiny calcifications, that help distinguish benign from malignant findings.
Instead, we use high-resolution patches focused on individual lesions. This preserves fine details, while the detection model provides the lesion's location. As a result, the classifier can focus on the most relevant part of the mammogram rather than processing the entire image.
The model
Any ImageNet backbone works. We swap the final layer for a single logit:
def build_classifier(arch="resnet50"):
if arch == "efficientnet_b0":
m = torchvision.models.efficientnet_b0(weights="DEFAULT")
m.classifier[1] = nn.Linear(m.classifier[1].in_features, 1)
else:
m = getattr(torchvision.models, arch)(weights="DEFAULT")
m.fc = nn.Linear(m.fc.in_features, 1)
return mPatches are normalized with ImageNet mean and standard deviation after repeating the grayscale channel three times, the same convention the pretrained weights were trained with.
Handling class imbalance
The dataset contains more benign lesions than malignant ones. Without adjustment, the model may favor predicting benign cases simply because they are more common.
To address this, we give malignant cases a higher weight during training. This weight is calculated as the number of benign cases divided by the number of malignant cases. As a result, missing a malignant lesion carries a greater penalty, helping the model pay more attention to the less common cancer cases:
n_pos = train_df.malignant.sum()
pos_weight = torch.tensor((len(train_df) - n_pos) / n_pos, device=device)
loss_fn = nn.BCEWithLogitsLoss(pos_weight=pos_weight)
opt = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)Training runs up to 40 epochs with cosine learning-rate decay and stops early after 8 epochs without validation AUROC improvement.
Choosing and freezing the threshold
The classifier produces a probability that a lesion is malignant. To turn this probability into a positive or negative prediction, we need to choose a threshold.
We select this threshold using the validation dataset and Youden's J statistic, which finds the best balance between sensitivity (detecting cancers) and specificity (correctly identifying benign lesions). The selected threshold is saved in the model checkpoint and used without changes during testing and inference:
from sklearn.metrics import roc_curve, roc_auc_score
y_val, p_val = predict(model, val_loader)
fpr, tpr, thr = roc_curve(y_val, p_val)
threshold = float(min(thr[np.argmax(tpr - fpr)], 1.0))
torch.save({"state_dict": model.state_dict(), "threshold": threshold}, "ckpt_classifier.pt")
y_test, p_test = predict(model, test_loader)
pred = p_test >= threshold # same threshold, never re-tuned
print("AUROC", roc_auc_score(y_test, p_test),
"sensitivity", (pred & (y_test == 1)).sum() / (y_test == 1).sum(),
"specificity", (~pred & (y_test == 0)).sum() / (y_test == 0).sum())
In a clinical screening system, the threshold would likely be chosen differently. For example, clinicians might require the model to detect at least 90% of cancers, even if this leads to more false positives.
We also plan to explore several improvements. Masses and calcifications have different visual characteristics, so training separate classifiers for each may produce better results than using a single model.
Another option is to evaluate each lesion patch in eight rotated and flipped orientations and average the predictions. This technique, known as test-time augmentation, may improve prediction stability.
Finally, we will report performance separately for masses and calcifications. A single overall AUROC score could look good while hiding poor performance on one type of lesion.
