Partie 2 : Amélioration avec la détection de ROI YOLOX

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

# 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

# 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}')

# 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)

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

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

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'].")

# 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).