J Regularization Improves Imbalanced Multiclass Segmentation

TL;DR

This paper introduces a novel regularization method using Youden's J statistic to improve multiclass cell segmentation, especially under class imbalance and weak supervision, achieving sharper boundaries and better separation.

Contribution

It proposes a new loss formulation with J statistic regularization that inherently handles class imbalance and enhances segmentation quality without explicit weighting.

Findings

Regularization improves boundary sharpness and cell separation.

Method performs well on 2D and 3D cell images with limited annotations.

Eliminates need for explicit class balancing weights.

Abstract

We propose a new loss formulation to further advance the multiclass segmentation of cluttered cells under weakly supervised conditions. We improve the separation of touching and immediate cells, obtaining sharp segmentation boundaries with high adequacy, when we add Youden's $J$ statistic regularization term to the cross entropy loss. This regularization intrinsically supports class imbalance thus eliminating the necessity of explicitly using weights to balance training. Simulations demonstrate this capability and show how the regularization leads to better results by helping advancing the optimization when cross entropy stalls. We build upon our previous work on multiclass segmentation by adding yet another training class representing gaps between adjacent cells. This addition helps the classifier identify narrow gaps as background and no longer as touching regions. We present results of our methods for 2D and 3D images, from bright field to confocal stacks containing different types of cells, and we show that they accurately segment individual cells after training with a limited number of annotated images, some of which are poorly annotated.