PK-YOLO: Pretrained Knowledge Guided YOLO for Brain Tumor Detection in Multiplanar MRI Slices

2024-10-29

IEEE Workshop/Winter Conference on Applications of Computer Vision

AI Summary

To address low detection accuracy for brain tumors—particularly small lesions (<32×32 pixels)—and difficulty in modeling cross-plane structural relationships in multi-planar MRI slices, this paper proposes the first knowledge-guided YOLO framework. Methodologically: (i) a lightweight CNN backbone is designed and pretrained via sparse-mask self-supervision to enhance MRI-specific feature representation; (ii) a knowledge transfer encoder is introduced to align features across orthogonal imaging planes; and (iii) an IoU-aware regression loss is formulated to improve sensitivity to small tumors. Evaluated on a multi-planar MRI brain tumor dataset, our method achieves a 4.2% absolute mAP improvement over state-of-the-art YOLO- and DETR-based detectors, with marked gains in small-tumor recall. The source code is publicly available.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Multi-instance/Multi-view Learning
• Intelligent Robots: Multimodal Perception & Sensor Fusion

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

Brain tumor detection in multiplane Magnetic Resonance Imaging (MRI) slices is a challenging task due to the various appearances and relationships in the structure of the multiplane images. In this paper, we propose a new You Only Look Once (YOLO)-based detection model that incorporates Pretrained Knowledge (PK), called PK-YOLO, to improve the performance for brain tumor detection in multiplane MRI slices. To our best knowledge, PK-YOLO is the first pretrained knowledge guided YOLO-based object detector. The main components of the new method are a pretrained pure lightweight convolutional neural network-based backbone via sparse masked modeling, a YOLO architecture with the pretrained backbone, and a regression loss function for improving small object detection. The pre-trained backbone allows for feature transferability of object queries on individual plane MRI slices into the model encoders, and the learned domain knowledge base can improve in-domain detection. The improved loss function can further boost detection performance on small-size brain tumors in multiplanar two-dimensional MRI slices. Experimental results show that the proposed PK-YOLO achieves competitive performance on the multiplanar MRI brain tumor detection datasets compared to state-of-the-art YOLO-like and DETR-like object detectors. The code is available at https://github.com/mkang315/PK-YOLO.

Problem

• Improving brain tumor detection in multiplanar MRI slices
• Enhancing small object detection with a new loss function
• Integrating pretrained knowledge into YOLO for better performance

Innovation

• Pretrained Knowledge guided YOLO architecture
• Lightweight CNN backbone via sparse masking
• Improved loss function for small tumors