YOLO-MED : Multi-Task Interaction Network for Biomedical Images

Backbone. The backbone network extracts features from input images. We choose CSPDarknet53 [8] as it supports feature propagation and reuse, leading to a significant reduction in parameter and computational overhead during training. This choice guarantees real-time network performance. The initial image data $\textbf{x}_{in}\in\mathbb{R}^{{H}{\times}W{\times}3}$ is input to the backbone.

Neck. The Neck is responsible for fusing the multi-scale features generated by the backbone. First, we pass the output of the backbone through an SPP (Spatial Pyramid Pooling) network [22] and subsequently feed it into the FPN (Feature Pyramid Network) [23]. The SPP module is utilized for feature generation and fusion across multiple scales, while the FPN module combines features from different semantic levels. This fusion process ensures that the resulting features encompass a rich blend of multi-scale and multi-semantic information. Within this module, we obtain three features with different scales: $\textbf{x}_{neck_{i}}\in\mathbb{R}^{\frac{H}{s_{i}}{\times}\frac{W}{s_{i}}{\times}c_{i}}$ where $s_{i}$ denotes the scale parameter ranging from 8 to 32, and $c_{i}$ represents the channel number of each feature map ranging from 128 to 512.

Decoupled heads for detection. First, we construct a Path Aggregation Network (PAN) [24]. PAN operates as a bottom-up pyramid network, aligning with the top-down semantic propagation in FPN. The diverse scale feature maps $\textbf{x}_{det_{i}}\in\mathbb{R}^{\frac{H}{s_{i}}{\times}\frac{W}{s_{i}}{\times}c_{i}}$ obtained from PAN are subsequently fused with the correspondingly scaled feature maps $\hat{\textbf{x}}_{det_{i}}\in\mathbb{R}^{\frac{H}{s_{i}}{\times}\frac{W}{s_{i}}{\times}c_{i}}$ generated by the cross-scale task-interaction module. These fused features serve as input to the final detection head. For our final detection head component, we opt for the decoupled head architecture [25]. This choice is based on the recognition that traditional coupled heads have demonstrated performance limitations due to inherent conflicts between classification and regression tasks [26]. Therefore, within our multi-task architecture, we introduce decoupled heads to ensure that each task (classification and regression) does not negatively affect the others.

Segment head. For the segmentation head, we devise a straightforward top-down network structure. Initially, we take the feature map from the lowest level of the FPN as input. Following two rounds of feature integration and upsampling, the resulting feature map $\textbf{x}_{seg}\in\mathbb{R}^{\frac{H}{2}{\times}\frac{W}{2}{\times}32}$ is merged with the feature map $\hat{\textbf{x}}_{seg}\in\mathbb{R}^{\frac{H}{2}{\times}\frac{W}{2}{\times}32}$ generated by the cross-scale task-interaction module. This fused feature map is subsequently used as input for the final round of feature integration and upsampling.

Implementation details. Our network is trained using Stochastic Gradient Descent (SGD) optimization algorithm with a learning rate of $1\times 10^{-2}$, weight decay of $5\times 10^{-4}$ and momentum of 0.937. Additionally, we employ the Cosine Annealing with Warm Restarts learning rate scheduling strategy, in which the first three epochs server as warm-up epochs with a reduced learning rate. To initiate the training, we utilize a pre-trained model from the COCO dataset. All experiments are conducted on a single NVIDIA GeForce 2080Ti GPU.

Baseline models and Metrics. We comprehensively evaluate our network by comparing it with various biomedical multi-task networks, as well as networks specialized in either object detection or semantic segmentation tasks. For object detection, we consider high-performing models from recent years in the biomedical image domain, including RetinaNet [9], as well as iconic models like Faster-RCNN [27] and YOLOv5s, representing two-stage and one-stage networks respectively. In addition, we include a comparison with our single-task baseline model, which consists of only the encoder and the detection decoder. We evaluate our model’s detection accuracy using mean Average Precision at 50% IoU (mAP50) and mean Average Precision at 95% IoU (mAP95) as metrics. Regarding the semantic segmentation task, our comparisons encompass classic architecture U-net [11], and Polyp-PVT [10] which utilizes transformer modules to enhance accuracy. We also include a comparison with the single-task baseline which comprises only the encoder and the segmentation decoder. We evaluate our model’s segmentation accuracy using Pixel Accuracy (PA) and mean Intersection over Union (meanIoU) as metrics. In the realm of multi-task networks, we compare our approach to the traditional UOLO [16], the latest MULAN [17], and the multi-task baseline of our model which includes only the encoder and two decoders. In addition to these horizontal comparisons with common networks, we conduct ablation experiments to investigate the impact of different modules within YOLO-Med, providing a detailed study of the network’s components.

Object detection results In the evaluation on the public dataset Kvasir-seg [21], as presented in Table 1, our model outperforms single-task networks such as Faster-RCNN [27], RetinaNet [9], YOLOv5s and our single-task baseline, as well as all three multi-task networks in terms of detection accuracy. Notably, our model demonstrates impressive real-time performance compared to other works, with only YOLOv5s surpassing it. However, it’s important to note that YOLOv5s lacks a segmentation decoder and a cross-scale task-interaction module. In the comparison between YOLO-Med and our single-task multi-task baseline models, all metrics indicate a higher level of object detection accuracy, with only a minimal decrease in inference speed. Similar results are observed on the private dataset as illustrated in Table 2.

Semantic segmentation results In the evaluation on the public dataset Kvasir-seg, as shown in Table 1, our network outperforms all three multi-task networks and three single-task networks. Furthermore, our network exhibit significantly superior real-time performance compared to both single-task or multi-task networks from other works, such as U-net [11], Polyp-PVT [10], UOLO [16] and MULAN [17]. When compared to our baseline models, all metrics surpass them, with an acceptable decrease in inference speed to enhance the segmentation accuracy. Similar results are observed on the large private dataset as shown in Table 2.