LabItem 15 of 22 · 2 hr

Part 2: Improving with YOLOX ROI Detection

What problem does YOLOX solve?

The threshold-based cropper from Part 1 fails on:

Implants (bright background)
Bright scanner markers/labels overlaid on the image
Low-contrast images where the breast edge is poorly defined
Cases where the background is not uniformly dark

The winning solution trains YOLOX-nano (a fast anchor-free object detector, 416×416 input) to directly predict the bounding box of the breast ROI. The crop is then resized to 2048×1024 for ConvNeXt.

Result: Cleaner, more consistent crops → improved ConvNeXt performance.

What you'll learn

Train a YOLOX-nano detector to localise the breast ROI bounding box on a mammogram
Build a YOLO-format detection dataset from pseudo-labels or manually annotated boxes
Wrap the trained detector to crop breast ROIs, with a threshold-crop fallback for missed detections
Regenerate the processed PNG dataset and retrain ConvNeXt on YOLOX-cropped images
Compare the YOLOX-cropped pipeline's pF1/AUC against the Part 1 threshold-crop baseline

Before you start

Completion of Part 1: Mammography Baseline — this notebook reuses df, CONFIG, read_dicom, and crop_breast_roi_threshold defined there
A GPU runtime (Kaggle T4/P100, or local CUDA) for both YOLOX-nano training and ConvNeXt retraining
Outbound network access to github.com, to clone the YOLOX repository
The RSNA dataset already downloaded, as in Part 1

Part 2.1 — Why YOLOX for Medical ROI Detection?

YOLOX is anchor-free and extremely fast at small sizes (nano = 0.91M params), making it ideal as a preprocessing step that must run on every image at inference time.

Aspect	Threshold cropper	YOLOX-nano
Speed	Very fast (CPU)	Fast (GPU, ~5ms)
Robustness	Fails on bright artefacts	Handles most cases
Training required	No	Yes (labelled boxes needed)
Generalisation	Scanner-dependent	Generalises across scanners

The winning team annotated 571 images manually (in YOLOv5 format) for training the detector.

Part 2.2 — YOLOX Setup


# Install YOLOX from the winning team's repo
!git clone https://github.com/Megvii-BaseDetection/YOLOX.git
%cd YOLOX
!pip install -v -e .   # install in editable mode
%cd ..

CheckpointYOLOX installed

You should see:


Successfully installed yolox-0.1.0

Not seeing this?

ModuleNotFoundError: No module named 'yolox' after pip install -v -e . — the cell ran from outside the YOLOX directory; re-run %cd YOLOX first.
CUDA-extension build error during install — the repo’s optional CUDA ops failed to compile; safe to ignore for CPU-only training/inference, the pure-Python fallback is used.
git clone hangs or times out — check outbound network access to github.com from the notebook environment.

Exercise 1 — Implement the pixel-to-YOLO and YOLO-to-pixel bounding-box conversion helpers.


# TODO
# Task A — Understand the annotation format:
#   YOLOX uses the YOLOv5 annotation format:
#   <class_id> <x_center> <y_center> <width> <height>  (all normalised 0-1)
#   For breast ROI there is only one class (class_id = 0 = breast).
#   Given a mammogram of shape (H=3000, W=1500), write a function that
#   converts a pixel bounding box (x1, y1, x2, y2) to this format.
 
def pixel_bbox_to_yolo(x1, y1, x2, y2, img_h, img_w):
    """
    Convert pixel (x1,y1,x2,y2) bbox to YOLO normalised format.
    Returns: class_id, x_center, y_center, width, height (all in [0,1])
    """
    # TODO: implement this
    raise NotImplementedError
 
 
def yolo_to_pixel_bbox(x_c, y_c, w, h, img_h, img_w):
    """
    Convert YOLO normalised format back to pixel (x1,y1,x2,y2).
    """
    # TODO: implement this
    raise NotImplementedError

Part 2.3 — Creating the ROI Detection Dataset

To train YOLOX we need bounding box annotations for breast ROIs.
Two options:

Use threshold cropper to generate pseudo-labels (quick, imperfect)
Download the winning team’s 571 manual annotations from the repo (better)


# Option 1: Auto-generate pseudo-labels from threshold cropper
# These will be noisy but sufficient for a reasonable detector.
 
import yaml
 
ROI_DATASET_DIR = './data/roi_det'
os.makedirs(f'{ROI_DATASET_DIR}/images/train', exist_ok=True)
os.makedirs(f'{ROI_DATASET_DIR}/images/val',   exist_ok=True)
os.makedirs(f'{ROI_DATASET_DIR}/labels/train', exist_ok=True)
os.makedirs(f'{ROI_DATASET_DIR}/labels/val',   exist_ok=True)
 
 
def generate_pseudo_label(row, src_dir, dst_img_dir, dst_lbl_dir):
    """Threshold-crop a DICOM, save resized PNG + YOLO annotation."""
    src = os.path.join(src_dir, str(row.patient_id), f'{row.image_id}.dcm')
    img = read_dicom(src)
    H, W = img.shape
 
    # Get bbox from threshold cropper
    mask = (img > 10).astype(np.uint8)
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
    if num_labels < 2:
        return None
    lbl = 1 + np.argmax(stats[1:, cv2.CC_STAT_AREA])
    x1 = stats[lbl, cv2.CC_STAT_LEFT]
    y1 = stats[lbl, cv2.CC_STAT_TOP]
    bw = stats[lbl, cv2.CC_STAT_WIDTH]
    bh = stats[lbl, cv2.CC_STAT_HEIGHT]
    x2, y2 = x1 + bw, y1 + bh
 
    # Save 416×416 resized image for YOLOX
    img_416 = cv2.resize(img, (416, 416))
    img_path = os.path.join(dst_img_dir, f'{row.patient_id}_{row.image_id}.png')
    cv2.imwrite(img_path, img_416)
 
    # Scale bbox to 416×416 and write YOLO label
    x1_s = x1 * 416 / W;  x2_s = x2 * 416 / W
    y1_s = y1 * 416 / H;  y2_s = y2 * 416 / H
    xc = (x1_s + x2_s) / 2 / 416
    yc = (y1_s + y2_s) / 2 / 416
    bw_n = (x2_s - x1_s) / 416
    bh_n = (y2_s - y1_s) / 416
 
    lbl_path = os.path.join(dst_lbl_dir, f'{row.patient_id}_{row.image_id}.txt')
    with open(lbl_path, 'w') as f:
        f.write(f'0 {xc:.6f} {yc:.6f} {bw_n:.6f} {bh_n:.6f}\n')
 
    return img_path
 
 
# Write dataset YAML for YOLOX
roi_yaml = {
    'path': ROI_DATASET_DIR,
    'train': 'images/train',
    'val':   'images/val',
    'nc': 1,
    'names': ['breast']
}
with open(f'{ROI_DATASET_DIR}/dataset.yaml', 'w') as f:
    yaml.dump(roi_yaml, f)
 
print("Dataset directory structure created.")
 
# Uncomment to run (slow — one DICOM per image):
# for _, row in tqdm(df.iterrows(), total=len(df)):
#     split = 'train' if random.random() > 0.1 else 'val'
#     generate_pseudo_label(row, CONFIG['train_images_dir'],
#                           f'{ROI_DATASET_DIR}/images/{split}',
#                           f'{ROI_DATASET_DIR}/labels/{split}')

CheckpointROI dataset directory populated

You should see:


Dataset directory structure created.
data/roi_det/
├── images/train
├── images/val
├── labels/train
├── labels/val
└── dataset.yaml

Not seeing this?

labels/train is empty after uncommenting the generation loop — generate_pseudo_label returned None for every row, meaning mask = (img > 10) found no connected component; check that read_dicom returns a non-empty array.
A label file contains nan values — img_h/img_w don’t match the DICOM’s actual shape; verify H, W = img.shape before scaling the bbox.

Part 2.4 — Training YOLOX-Nano

Exercise 2 — Write the YOLOX-nano experiment file and launch detector training.


# TODO
# Task A — Experiment file:
#   YOLOX uses Python experiment files (exps/) to configure training.
#   Create exps/rsna_yolox_nano.py based on the nano template,
#   setting num_classes=1, input_size=(416,416), max_epoch=50.
#
# Task B — Run training:
#   python YOLOX/tools/train.py -f exps/rsna_yolox_nano.py -d 1 -b 16 --fp16
#   Monitor mAP@0.5 on the val split. The winning team reports ~95% AP@0.5.
#
# Task C — Why nano and not a larger YOLOX?
#   The ROI detection task is simple (one large object per image, near-perfect
#   contrast). A nano model (0.91M params) is sufficient and runs fast.
#   Verify: does a larger YOLOX-s actually improve downstream ConvNeXt pF1?
 
# Example training command (run in terminal):
YOLOX_TRAIN_CMD = """
PYTHONPATH=$(pwd)/YOLOX:$PYTHONPATH python YOLOX/tools/train.py \\
    -f exps/rsna_yolox_nano.py \\
    -d 1 \\
    -b 16 \\
    --fp16 \\
    -o \\
    --cache
"""
print("Training command:")
print(YOLOX_TRAIN_CMD)

CheckpointYOLOX-nano training launched

You should see:


Training command:
PYTHONPATH=$(pwd)/YOLOX:$PYTHONPATH python YOLOX/tools/train.py \
    -f exps/rsna_yolox_nano.py \
    -d 1 \
    -b 16 \
    --fp16 \
    -o \
    --cache
...
mAP@0.5 ~ 0.95 on val split

Not seeing this?

CUDA out of memory — lower -b (batch size) from 16 to 8; nano is small but shared-GPU notebooks can still OOM.
mAP stuck near 0 after 50 epochs — num_classes in the experiment file doesn’t match the single breast class, or dataset.yaml’s path/train/val fields point to the wrong directory.
--cache fails with a disk-space error — pass --cache ram or drop the flag; caching resized images to disk needs scratch space for the full training set.

Part 2.5 — YOLOX Inference for ROI Cropping

Define a wrapper that loads a trained YOLOX-nano and returns a breast bounding box on the original image scale.


import sys
sys.path.insert(0, 'YOLOX')
 
from yolox.data.data_augment import ValTransform
from yolox.data.datasets import COCO_CLASSES
from yolox.exp import get_exp
from yolox.utils import fuse_model, get_model_info, postprocess
 
 
class YOLOXBreastDetector:
    """
    Wrapper around a trained YOLOX-nano model for breast ROI detection.
    Produces a (x1, y1, x2, y2) bounding box on the original image scale.
    """
    def __init__(self, exp_file: str, ckpt_path: str, device: str = 'cuda',
                 input_size: tuple = (416, 416), score_thresh: float = 0.3):
        self.input_size   = input_size
        self.score_thresh = score_thresh
        self.device       = device
 
        exp = get_exp(exp_file, None)
        exp.test_size = input_size
 
        self.model = exp.get_model()
        ckpt = torch.load(ckpt_path, map_location=device)
        self.model.load_state_dict(ckpt.get('model', ckpt))
        self.model = fuse_model(self.model).to(device).eval()
 
        self.preproc = ValTransform(legacy=False)
 
    @torch.inference_mode()
    def detect(self, img_gray: np.ndarray):
        """
        Args:
            img_gray: uint8 grayscale mammogram array (H, W)
        Returns:
            bbox (x1, y1, x2, y2) in original image pixels, or None if no detection
        """
        H, W = img_gray.shape
        img_rgb = cv2.cvtColor(img_gray, cv2.COLOR_GRAY2RGB)
 
        # Preprocess to YOLOX input size
        img_t, ratio = self.preproc(img_rgb, None, self.input_size)
        img_t = torch.from_numpy(img_t).unsqueeze(0).float().to(self.device)
 
        # Run YOLOX
        outputs = self.model(img_t)
        outputs = postprocess(outputs, num_classes=1, conf_thre=self.score_thresh,
                              nms_thre=0.45, class_agnostic=True)
 
        if outputs[0] is None or len(outputs[0]) == 0:
            return None  # no detection — fall back to threshold crop
 
        # Take highest-confidence detection
        boxes = outputs[0].cpu().numpy()
        best  = boxes[np.argmax(boxes[:, 4])]
        x1, y1, x2, y2 = best[:4] / ratio
 
        # Clamp to image bounds
        x1 = max(0, int(x1));  y1 = max(0, int(y1))
        x2 = min(W, int(x2));  y2 = min(H, int(y2))
        return x1, y1, x2, y2
 
 
print("YOLOXBreastDetector class defined.")
print("Instantiate with:")
print("  detector = YOLOXBreastDetector(")
print("      exp_file='exps/rsna_yolox_nano.py',")
print("      ckpt_path='YOLOX/YOLOX_outputs/rsna_yolox_nano/best_ckpt.pth'")
print("  )")


def crop_with_yolox(img_gray: np.ndarray, detector: YOLOXBreastDetector,
                    fallback_threshold: bool = True) -> np.ndarray:
    """
    Crop breast ROI using YOLOX. Falls back to threshold cropping if no
    detection is found (robustness measure).
    """
    bbox = detector.detect(img_gray)
    if bbox is not None:
        x1, y1, x2, y2 = bbox
        return img_gray[y1:y2, x1:x2]
    elif fallback_threshold:
        return crop_breast_roi_threshold(img_gray)
    else:
        return img_gray

CheckpointDetector wrapper returns a valid bbox

You should see:


YOLOXBreastDetector class defined.
Instantiate with:
  detector = YOLOXBreastDetector(
      exp_file='exps/rsna_yolox_nano.py',
      ckpt_path='YOLOX/YOLOX_outputs/rsna_yolox_nano/best_ckpt.pth'
  )
>>> detector.detect(img_gray)
(142, 88, 1390, 2872)

Not seeing this?

detect() always returns None — score_thresh=0.3 may be too high for an undertrained checkpoint; lower it temporarily to confirm the model produces any detections at all.
RuntimeError: size mismatch in load_state_dict — ckpt_path points at a checkpoint trained with a different experiment file (e.g. yolox-s instead of nano); re-check the -f flag used at training time.
Returned bbox has x2 <= x1 or y2 <= y1 — confirm img_gray.shape is (H, W), not (W, H).

Exercise 3 — Compare threshold vs YOLOX crops, measure detection coverage, and analyse confidence scores.


# TODO
# Task A — Compare crop quality side by side:
#   For 6 images (2 normal, 2 with artefacts, 2 implants):
#   Show: original | threshold crop | YOLOX crop
#   Mark the predicted bounding box on the original image.
#
# Task B — Measure coverage:
#   Compute what fraction of images YOLOX successfully detects vs falls back
#   to threshold cropping. What are the characteristics of failed detections?
#
# Task C — YOLOX confidence analysis:
#   Plot the distribution of detection confidence scores.
#   Do low-confidence detections produce worse crops?
#   Consider using a higher score_thresh (e.g. 0.5) and more aggressive fallback.

Part 2.6 — Regenerate Processed PNGs with YOLOX Crops

Now rerun the DICOM→PNG conversion pipeline from Section 3, but replace crop_breast_roi_threshold with crop_with_yolox.


PROCESSED_YOLOX_DIR = './data/processed_pngs_yolox'
os.makedirs(PROCESSED_YOLOX_DIR, exist_ok=True)
 
 
def convert_dicom_to_png_yolox(row, src_dir: str, dst_dir: str, detector):
    """DICOM → 8-bit normalise → YOLOX crop → PNG."""
    src = os.path.join(src_dir, str(row.patient_id), f'{row.image_id}.dcm')
    dst = os.path.join(dst_dir, f'{row.patient_id}_{row.image_id}.png')
 
    if os.path.exists(dst):
        return dst
 
    img = read_dicom(src)
    img = crop_with_yolox(img, detector, fallback_threshold=True)
    cv2.imwrite(dst, img)
    return dst
 
 
# Uncomment after training YOLOX:
# detector = YOLOXBreastDetector(
#     exp_file='exps/rsna_yolox_nano.py',
#     ckpt_path='YOLOX/YOLOX_outputs/rsna_yolox_nano/best_ckpt.pth'
# )
# for _, row in tqdm(df.iterrows(), total=len(df)):
#     convert_dicom_to_png_yolox(row, CONFIG['train_images_dir'], PROCESSED_YOLOX_DIR, detector)
 
# Update paths in df
# df['path'] = df.apply(
#     lambda r: os.path.join(PROCESSED_YOLOX_DIR, f"{r.patient_id}_{r.image_id}.png"), axis=1
# )
print("After regenerating PNGs, rerun Section 9 (CV training) with the updated df['path'].")

CheckpointYOLOX-cropped PNGs regenerated

You should see:


>>> len(os.listdir(PROCESSED_YOLOX_DIR)) == len(df)
True
After regenerating PNGs, rerun Section 9 (CV training) with the updated df['path'].

Not seeing this?

PROCESSED_YOLOX_DIR has fewer PNGs than rows in df — some detect() calls fell back to crop_breast_roi_threshold, which can still raise on a fully empty mask; log and skip failing image_ids instead of aborting the loop.
Re-running the loop appears to do nothing — convert_dicom_to_png_yolox short-circuits with if os.path.exists(dst): return dst; delete PROCESSED_YOLOX_DIR first if you changed score_thresh and want to regenerate.

Part 2.7 — Retrain ConvNeXt with YOLOX-Cropped Images

Exercise 4 — Retrain the 4-fold ConvNeXt on YOLOX crops and compare pF1 against the baseline.


# TODO
# Task A — Retrain and compare:
#   Run the full 4-fold CV from Section 9 again, but with
#   YOLOX-cropped images (df['path'] pointing to PROCESSED_YOLOX_DIR).
#   Fill in the table below:
#
#   | Crop method   | OOF pF1 | OOF AUC |
#   |---            |---      |---      |
#   | Threshold     | ?       | ?       |
#   | YOLOX-nano    | ?       | ?       |
#
# Task B — Error analysis on improved crops:
#   Identify images where YOLOX cropping changed the prediction significantly
#   (|pred_yolox - pred_threshold| > 0.2). Are these the artefact/implant cases?
#
# Task C — Larger YOLOX vs YOLOX-nano:
#   Try training YOLOX-s (small, 9M params). Does the better detection
#   quality translate to better ConvNeXt pF1? Or is YOLOX-nano already
#   good enough (the winning answer from the mr.robot team is: nano is sufficient).

Clean up

The YOLOX checkpoints, regenerated PNGs, and retrained ConvNeXt fold weights are yours to keep. Release the compute once you’re done comparing pipelines:

On Kaggle, stop the session (or switch the accelerator back to None) after the YOLOX training run and the ConvNeXt retraining run both finish, to free the GPU and stop drawing down your weekly quota.
On a local workstation, shut down the notebook kernel to release GPU memory once the pF1/AUC comparison table is filled in.

Summary: From Baseline to Winning Pipeline

Pipeline	Expected OOF pF1
Part 1 — Baseline	~0.52–0.56
Part 2 — Full winning pipeline	~0.59–0.62 (LB 0.65, AUC 0.93 with ensemble)

Further improvements the winning team explored (but are out of scope here):

External data (VinDr, CMMD, CBIS-DDSM) for backbone pretraining
TTA (horizontal flip ensemble at inference)
All 4 views (CC+MLO, L+R) as a patient-level prediction
MONOCHROME1 inversion verified per-scanner
MaxPool head instead of AvgPool (already implemented above)

Reference:
mr.robot team writeup: https://www.kaggle.com/competitions/rsna-breast-cancer-detection/writeups/mr-robot-1st-place-solution
Code: https://github.com/dangnh0611/kaggle_rsna_breast_cancer