\section{Method}
\label{sec:method}

\subsection{Problem Formulation}
\label{subsec:problem}
This study aims to segment an array of anatomical concepts in renal pathology, encapsulating three conceptual layers — regions ($R_1, R_2$), functional units ($F_1, F_2, F_3, F_4$), and cells ($C_1, C_2$) — spanning 8 distinct objects. 


Leveraging anatomical learning, our approach is tailored to achieve comprehensive segmentation in renal pathology by effectively interpreting both affirmative and negatory relationships within the anatomical relationship.

\begin{figure*}[t]
\centering 
\includegraphics[width=1\linewidth]{{fig/Architecture.pdf}}
\caption{
\textbf{Token-based dynamic head network architecture -- }
This figure illustrates the architecture of the proposed PrPSeg method. It incorporates a residual U-Net backbone, augmented with class-aware and scale-aware tokens. These tokens are integrated into each block of the encoder, as well as the Global Average Pooling (GAP) block, ensuring a comprehensive understanding of both class and scale features. Such features are aggregated by a fusion block to adaptively generate the parameters for a single dynamic segmentation head. The proposed method is able to segment all hierarchical semantic anatomies using a single network.} 
\label{Fig.architecture} 
\end{figure*}


The pipeline is composed of three integral components:

\noindent\textbf{Universal proposition matrix:} A meticulously crafted universal proposition matrix is employed to elucidate the anatomical relationships among various objects. This matrix aids in enhancing structural understanding from an engineering perspective, facilitating better cognition of the complex renal architecture.

\noindent\textbf{Token-based dynamic head backbone:} We have developed a token-based dynamic head backbone architecture that is adept at interpreting class-aware and scale-aware knowledge pertinent to renal pathology. This component is pivotal in accommodating the extensibility requirements posed by the introduction of new data, ensuring the model's adaptability and scalability.

\noindent\textbf{Anatomy loss function:} A novel anatomy loss function has been formulated, which operationalizes the affirmative and negatory relationships inherent in kidney anatomy. This function is a critical element in achieving nuanced, all-encompassing segmentation, bolstering the model's ability to discern and categorize the intricate structures within the kidney.

\subsection{Universal proposition matrix}
\label{subsec:matrix}
Renal pathology encompasses regions (the medulla and cortex), functional units (glomerulus, tubules, etc.), with corresponding cellular structures. Pathologists analyze the morphology of the kidney by examining the functional flow through these heterostructures. Traditionally, each heterostructure undergoes isolated examination and quantification to meet specific demands, often overlooking the homostructures within each unit, such as podocyte cells, mesangial cells in glomerular tufts. To enhance the understanding of the integrated kidney structure, we propose transforming kidney anatomy into an anatomical relationship map, using principles of affirmation and negation from linguistic and grammatical concepts.

To translate these anatomy concepts into an engineering framework, we adopt Aristotle's logic theory to develop an anatomical map for renal pathology. Aristotle's theory examines the relationships between objects using four fundamental categorical propositions. Upon closer inspection, complex propositions reveal themselves as collections of simpler claims derived from these initial propositions. Specifically, we utilize two terms from Aristotle's theory: \textit{(1) Universal Affirmation: ``All S are P," and (2) Universal Negation: ``No S are P,"}, to universally assert properties for all group members, indicating strong constraints and relationships. For example, in the context of two kidney structures, $A$ and $B$, if $B$ is within $A$, $B$ is a subset of $A$, following the rule of Universal Affirmation. Conversely, if $A$ and $B$ have no inclusion relationship, they are mutually exclusive, aligning with Universal Negation. These propositions are employed to construct an anatomical relationship map representing the classification and spatial relationships among renal pathology objects, as illustrated in~\Cref{fig.anatomy}.

This anatomical relationship is characterized by:

\noindent\textbf{Uniqueness:} Each pair of objects is linked by a single proposition. The expanding structure of the map, devoid of cycles, ensures stable inheritance relationships from regions down to cellular levels.

\noindent\textbf{Transmissibility:} Indirect relationships between objects can be deduced from direct relationships, as established by the two fundamental categorical propositions. Relationships between objects not directly connected can be determined by combining propositions along their connecting path.

Following the translation of this knowledge from clinical anatomy to an engineering paradigm, we introduce a universal proposition matrix, $M_t\in\mathbb{R}^{n \times n}$, to facilitate implementation in computational models. Here, $n$ represents the number of classes within the map. The matrix values are defined in~\Cref{eq:matrix}:


\vspace{-5mm} 
\begin{equation}
    M_t(i,j) = 
    \begin{cases}
    1, \quad if I  \subseteq J \\
    -1, \quad if I \supseteq J \\
    2, \quad if I \cap J = \emptyset \\
    0, \quad \text{otherwise}
    \end{cases} 
    i,j = 1,2,...,n
\label{eq:matrix}
\end{equation}

This matrix is designed to be extendable with the introduction of new data.


\subsection{Token-based Dynamic Head Network}
\label{subsec:token}


\begin{figure}
\centering 
\includegraphics[width=1\linewidth]{{fig/Backbone.pdf}}
\caption{
\textbf{Token-based dynamic head -- }
This figure visualizes the architecture of our proposed token-based dynamic head backbone. Central to our design is the ability to maintain a consistent model architecture while dynamically accommodating an increasing number of segmentation classes. This flexibility is achieved by extending the dimensions of the tokens, rather than altering the backbone structure. Key components include a dynamic token bank with class-aware and scale-aware tokens, an encoder, and a dynamic head network, all orchestrated to efficiently handle class expansion without necessitating changes to the backbone.} 
\label{Fig.token} 
\end{figure}

Pathological image segmentation faces three main challenges: (1) Heterogeneous object annotations are often partially labeled, with only one type of tissue annotated per pathological image; (2) It is challenging to form anatomical relationships (e.g. subset/superset relationships) in multi-class segmentation; for example, it is difficult to simultaneously segment glomerular capsule, tufts, and cells that share regions; (3) The annotation process on giga-pixel images is labor-intensive, leading to an ongoing data collection process with class extension. Therefore, a segmentation backbone optimized for binary segmentation of multiple classes with spatial overlap, and adaptable to data enlargement, is required. Previous designs, including multi-head and dynamic-head, are suboptimal for partially labeled learning and data extension, leading to changes in the backbone architecture and insufficient performance.

In this work, we propose a token-based dynamic head backbone designed to maintain consistent model architecture while accommodating a possibly increasing number of segmentation classes (In~\Cref{Fig.token}). The backbone can be supervised with partially labeled data, understanding anatomical relationships effectively. The entire architecture is demonstrated in~\Cref{Fig.architecture}. The backbone of our proposed method is the Residual-U-Net from Omni-Seg~\cite{deng2023omni}, chosen for its superior segmentation performance in renal pathology. Instead of using dimensionally-changeable one-hot vectors for class-aware encoding and scale-aware encoding, we use dimensionally-stable class-aware tokens ($T_c\in\mathbb{R}^{n \times d}$) and scale-aware tokens ($T_s\in\mathbb{R}^{4 \times d}$) to comprehend the contextual information in the encoder ($E$) of the model. Here, $d$ is the sum of the channel numbers of blocks ($d_1+ d_2+\dots+d_b+d_{gap}$) in the encoder. Each interval of the channels represents level-specific features in each level of the encoder. Each class has a one-dimensional token $t_c\in\mathbb{R}^{1 \times d}$ to store class-specific knowledge at the feature level among the whole dataset, while each magnification has a one-dimensional token $t_s\in\mathbb{R}^{1 \times d}$ for providing scale-specific knowledge across four scales (5$\times$, 10$\times$, 20$\times$, and 40$\times$). Inspired by the Vision Transformer~\cite{dosovitskiy2020image}, for an image $I$ of class $i$ with magnification $m$, at the $b$-th encoder block, the corresponding class token $T_c(i)$ and $T_s(m)$ are added to the current feature map ($e_{b-1}$) before being fed into the current block ($E_b$). This process is defined by the following equation:

\begin{equation}
    \text{Ind}_{start} = \sum_{k=1}^{b-1} d_k,  \text{Ind}_{end} = \sum_{k=1}^{b} d_k
\label{eq:fusion}
\end{equation}

\begin{equation}
\begin{aligned}
    e_b = E_b(&  T_c[i][\text{Ind}_{start}:\text{Ind}_{end}] + \\
    &   T_s[m][\text{Ind}_{start}:\text{Ind}_{end}] + e_{b-1})
\end{aligned}
\label{eq:token}
\end{equation}

With the class token $T_c$ and scale token $T_s$, the encoder captures domain-specific features in the image.

To further integrate class-specific and scale-specific information into the embedded features, we combine the last intervals of the class-token and scale token into the low-dimensional feature embedding at the bottom of the Residual-U-Net architecture. The image feature $F$ is summarized by Global Average Pooling (GAP) and transformed into a feature vector in $\mathbb{R}^{d_{gap}}$. $d_{gap}$ is the dimension of GAP feature. Differing from Omni-Seg~\cite{deng2023omni}, which implements a triple outer product to combine three vectors into a one-dimensional vector via a flatten function, we use a single 2D convolutional layer controller, $\varphi$, as a feature fusion block to refine the fusion vector as the final controller for dynamic head mapping:


\vspace{-3mm}
\begin{equation}
    \omega = \varphi(\text{GAP}(F)+T_c[i][:-d_{gap}]+T_s[i][:-d_{gap}];\Theta_\varphi)
\label{eq:fusion}
\end{equation}

\noindent where $\text{GAP}(F)$, $T_c$, and $T_s$ are combined by the addition operation, and $\Theta_\varphi$ represents the number of parameters in the dynamic head.

Following the approach of~\cite{deng2023omni}, a binary segmentation network is used to achieve multi-label segmentation via a dynamic filter. From the multi-label multi-scale modeling described above, we derive joint low-dimensional image feature vectors, class-specific tokens, and scale-specific tokens at an optimal segmentation magnification. These are then mapped to control a lightweight dynamic head, specifying (1) the target tissue type and (2) the corresponding pyramid scale.

The dynamic head consists of three layers. The first two layers contain eight channels each, while the final layer comprises two channels. We directly map parameters from the fusion-based feature controller to the kernels in the 162-parameter dynamic head to achieve precise segmentation from multi-modal features. The filtering process is expressed by the following equation:


\vspace{-2mm}
\begin{equation}
    P = ((((M \cdot \omega_1) \cdot \omega_2) \cdot \omega_3)
\label{eq:dynamic-head}
\end{equation}


\noindent where $\cdot$ denotes convolution, $P\in\mathbb{R}^{2\times W \times H}$ is the final prediction, and $W$ and $H$ correspond to the width and height of the dataset, respectively.

The benefits of the dynamic-token design are twofold: (1) The backbone architecture remains stable and reusable as new classes are introduced, which is advantageous for incremental learning. (2) The binary segmentation scheme allows the model to predict multiple classes with spatial overlap, outperforming other multi-head and dynamic-head designs.

\subsection{Anatomy Loss Function}
\label{subsec:loss}

With the introduction of the universal proposition matrix and token-based dynamic head architecture, we propose an online semi-supervised anatomy learning strategy to incorporate spatial correlation into the training process for comprehensive segmentation. For a given image $I$ with a labeled class $i$, represented as $Y_i$, we generate predictions $Y'_j$ for another class $j$ on the same image. We then use the anatomical relationship to supervise the correlation between the supervised label $Y_i$ and the semi-supervised prediction $Y'_j$: (1) If $i$ is a superset of $j$, $Y'_j$ should not exceed the region of $Y_i$; conversely, (2) if $i$ is a subset of $j$, $Y'_j$ should cover $Y_i$ as comprehensively as possible; and (3) if $i$ is mutually exclusive with $j$, the overlap between $Y_i$ and $Y'_j$ should be minimized. The total anatomy loss is defined by the following equations:

% \begin{equation}
% m = M(i,j)
% \label{eq:matrix_value}
% \end{equation}

% \begin{equation}
% \begin{aligned}
%     L_{\text{hats}}(i,j) =  & \frac{(m+1)(m-2)}{-2} \cdot \text{DCE}(1 - Y_i, Y'_j)\\
%     & - \frac{(m-1)(m-2)}{6} \cdot \text{DCE}(Y_i, Y_i \cup Y'_j ) \\
%     & + \frac{(m-1)(m+1)}{3} \cdot \text{DCE}(Y_i, Y'_j) \\
% \end{aligned}
% \label{eq:anatomy_loss}
% \end{equation}

\vspace{-5mm}
\begin{equation}
    L_{\text{upl}}(i,j) = 
    \begin{cases}
    \text{DCE}(1 - Y_i, Y'_j), \quad \text{if} \quad M_t(i,j) = 1 \\
    - \text{DCE}(Y_i, Y_i \cup Y'_j ), \quad \text{if} \quad M_t(i,j) = -1 \\
    \text{DCE}(Y_i, Y'_j), \quad \text{if} \quad M_t(i,j) = 2 \\
    0, \quad \text{if} \quad M_t(i,j) = 0 \\
    \end{cases} 
\label{eq:anatomy_loss}
\end{equation}
\vspace{-3mm}

\noindent where $M_t$ is the anatomy matrix and $\text{DCE}$ denotes the Dice Loss. The total loss function is an aggregate of supervised and semi-supervised losses, weighted by $\lambda_{\text{upl}}$.

\vspace{-2mm}
\begin{equation}
\begin{aligned}
    L(i) = & \text{DCE}(Y_i,Y'_i) + \text{BCE}(Y_i,Y'_i)\\
    & + \lambda_{\text{upl}} \sum_{j=1}^{n}L_{\text{upl}}(i,j) \quad ( j \neq i )
\end{aligned}
\label{eq:loss_function}
\end{equation}
\vspace{-5mm}

\noindent where $\text{BCE}$ represents the Binary Cross-Entropy loss, and $Y'_i$ is the prediction for class $i$.
