Class-balanced Open-set Semi-supervised Object Detection for Medical Images

TL;DR

The paper tackles open-set semi-supervised object detection (OSSOD) in medical imaging, where unlabeled data may contain unknown disease categories and class imbalance is severe. It introduces two modules, Out-of-distribution Detection Fusion Classifier (OODFC) and Category Control Embed (CCE), integrated into a Mean-Teacher SSOD framework to mitigate unknown-category interference and balance foreground categories. OODFC merges unknown-class pseudo-labels with known ones using a dynamic threshold based on category AP, while CCE constructs a Foreground Information Library and synthesizes balanced foreground data via $x_i^{Syn}=\beta b_j^{'}+(1-\beta)x_i$, $y_i^{Syn}=y_j$, with $\beta=0.5$, all within an EMA-based training regime. Experiments on private ophthalmology and public Parasite datasets show consistent mAP gains, including a 4.25-point improvement on Parasite, demonstrating effective open-set handling and class balance with practical medical-imaging impact.

Abstract

Medical image datasets in the real world are often unlabeled and imbalanced, and Semi-Supervised Object Detection (SSOD) can utilize unlabeled data to improve an object detector. However, existing approaches predominantly assumed that the unlabeled data and test data do not contain out-of-distribution (OOD) classes. The few open-set semi-supervised object detection methods have two weaknesses: first, the class imbalance is not considered; second, the OOD instances are distinguished and simply discarded during pseudo-labeling. In this paper, we consider the open-set semi-supervised object detection problem which leverages unlabeled data that contain OOD classes to improve object detection for medical images. Our study incorporates two key innovations: Category Control Embed (CCE) and out-of-distribution Detection Fusion Classifier (OODFC). CCE is designed to tackle dataset imbalance by constructing a Foreground information Library, while OODFC tackles open-set challenges by integrating the ``unknown'' information into basic pseudo-labels. Our method outperforms the state-of-the-art SSOD performance, achieving a 4.25 mAP improvement on the public Parasite dataset.

Figures (3)

• Figure 1: This is a visualization of the prediction results of three models on four images that contain unknown objects. In the predictions of soft-teacher and DSL, the unknown categories are incorrectly predicted as the Ancylostoma Spp category, which appears most frequently in the training set. In contrast, OODFC can produce "unknown" labels.
• Figure 2: Overall structure of our framework. We designed two modules to address the class imbalance and open-set problems. The entire framework is based on a semi-supervised pipeline using a teacher-student network. Weak augmentation includes random flip, while strong augmentation includes random flip, color jittering, and cutout. The CCE mitigates category imbalance by dynamically embedding foreground information, while the OODFC prevents the model from misclassifying unknown categories as known ones by integrating unknown class information.
• Figure 3: Visualization of synthetic images generated by CCE. In these examples, several foreground image segments are randomly added to the original unlabeled image.