Improving Diagnostic Performance on Small and Imbalanced Datasets Using Class-Based Input Image Composition

TL;DR

This work tackles the challenge of high diagnostic error rates on small, imbalanced medical imaging datasets by introducing Class-Based Input Image Composition (CB-ImgComp). CB-ImgComp constructs Composite Input Images (Co-Img) in structured layouts (notably $3\times1$) by merging multiple images from the same class, thereby increasing intra-class variability and information density. The approach enables perfect or near-perfect balancing and substantial performance gains, demonstrated by a VGG16 model achieving approximately $99.7\%$ accuracy and $AUC\approx0.9996$ on the Co-OCTDL dataset derived from OCTDL. The findings suggest that input-level data representation and augmentation can outperform purely architectural improvements, offering a practical, model-agnostic path to robust diagnostics in data-scarce clinical settings.

Abstract

Small, imbalanced datasets and poor input image quality can lead to high false predictions rates with deep learning models. This paper introduces Class-Based Image Composition, an approach that allows us to reformulate training inputs through a fusion of multiple images of the same class into combined visual composites, named Composite Input Images (CoImg). That enhances the intra-class variance and improves the valuable information density per training sample and increases the ability of the model to distinguish between subtle disease patterns. Our method was evaluated on the Optical Coherence Tomography Dataset for Image-Based Deep Learning Methods (OCTDL) (Kulyabin et al., 2024), which contains 2,064 high-resolution optical coherence tomography (OCT) scans of the human retina, representing seven distinct diseases with a significant class imbalance. We constructed a perfectly class-balanced version of this dataset, named Co-OCTDL, where each scan is resented as a 3x1 layout composite image. To assess the effectiveness of this new representation, we conducted a comparative analysis between the original dataset and its variant using a VGG16 model. A fair comparison was ensured by utilizing the identical model architecture and hyperparameters for all experiments. The proposed approach markedly improved diagnostic results.The enhanced Dataset achieved near-perfect accuracy (99.6%) with F1-score (0.995) and AUC (0.9996), compared to a baseline model trained on raw dataset. The false prediction rate was also significantly lower, this demonstrates that the method can producehigh-quality predictions even for weak datasets affected by class imbalance or small sample size.

Figures (6)

• Figure 1: Constructing enhanced input images by merging same-class samples.
• Figure 2: Co-Dataset Generator Algorithm Schema.
• Figure 3: Co-OCDTL AMD Class examples
• Figure 4: Dataset size expansion using CB-ImgComp generation with $k$=3
• Figure 5: Comparison of Initial Dataset and Balanced Co-Dataset