RevisedMedYOLO: Unlocking Model Performance by Careful Training Code Inspection

Kai Geissler, Jan Hendrik Moltz, Hans Meine, Markus Wenzel

Keywords: MedYOLO, Medical Object Detection, Lesion Detection, 3D Object Detection

TL;DR: We introduce RevisedMedYOLO, a corrected version of MedYOLO for 3D medical object detection, achieved by fixing critical bugs in the original training code, enabling successful training on LIDC and BraTS where the original model previously failed.

Abstract: The MedYOLO architecture, adapting YOLOv5 for 3D medical object detection, was reported by its original authors to have failed to train effectively in LIDC for the detection of lung lesions and to fail in BraTS for the detection of brain tumors when using a large model configuration. This work introduces RevisedMedYOLO, achieved by carefully reviewing and correcting the original training implementation. We fixed critical bugs related to dataset shuffling, initialization of bias parameters, and bounding box clamping during zoom augmentation. Consequently, RevisedMedYOLO demonstrates successful learning on these datasets (AP@$0.5>0$), unlike the original implementation. This study underscores the crucial role of careful code implementation and debugging in enabling deep learning model performance for challenging tasks in medical image analysis.

Submission Number: 56