CountMamba: Exploring Multi-directional Selective State-Space Models for Plant Counting

Counting the number of plants significantly predicts crop yield and other aspects. Initially, the counting of plants is usually distinguished from other backgrounds through different colors [21] in the agricultural field, as there exists a strong color contrast between different plants. Nevertheless, this method is difficult to distinguish when plants are obstructed or have similar colors. Based on this issue, many methods attempt to solve the problem through phenotypic parameters such as plant appearance texture features [5]. However, distinguishing phenotypic parameters was also difficult at the time. With the emergence of deep learning methods, the difficulty is mainly alleviated.

In deep learning, many methods tend to construct segmentation [6] or detection [19, 24, 35] pipelines to locate plants. Because both detection and segmentation have well-known, powerful, and easy-to-use frameworks, such as faster R-CNN [27] and fully convolutional networks (FCN) [20], apart from the traditional convolutional networks, new network architectures such as Transformer [30] and Mamba [9] that have emerged in recent years can also be utilized for extracting plant phenotypic characteristics. Harada et al. [28] exploit a novel hybrid wheat detection model by incorporating the CNN and Transformer for modeling long-range dependence. WheatNet [35] consists of two parts: a Transform Network, which reduces the effect of differences in the color features, and a Detection Network, which is designed to improve the capability of detection. FlowerNet [16] is a framework for flower counting based on the algorithm of YOLACT++ [3], a real-time instance segmentation model. TasselNet [22] is the first work of a local count network applied to plants. TasselNetV2 [33] discovers weak context on the plants, which is essential for the counting work.

Apart from the segmentation and detection abundantly utilized in plant counting, regression-based plant counting is less adopted. Giuffrida et al. [8] introduce scale and rotation invariance to learn features in a log-polar representation, which aims to count leaves in round shapes. A combined network that integrates density map regression and image segmentation is proposed by Wu et al. [32] for rice seedlings counting. Regression-based methods generally use the overall features of the image for regression analysis, ignoring the spatial information of the image. Therefore, their performance is worse than that of some density map-based methods constructed by detection or segmentation pipelines.

State Space Models (SSMs) have shown remarkable and efficient ability in capturing long sequences by utilizing selective state space to capture relevant information. Unlike Transformer [30], SSMs have lower linear time complexity and implement linear time running in terms of sequence length, making them particularly suitable for handling very long sequences. Gu et al. [10] first introduce a Structured State-Space Sequence model (S4) engineered with a specific focus on capturing long-range dependencies, which boasts the advantage of linear complexity. Inspired by this, various models have been developed, such as the S5 layer [29], which introduces MIMO SSM and efficient parallel scan into the S4 layer. Fu et al. [7] propose a novel SSM layer, H3, which significantly narrows the performance disparity between SSMs and Transformer attention in language modeling. Methea et al. [25] enhance the expressiveness of the Gated State Space layer on S4 by incorporating additional gating units, thereby boosting its performance. Lately, Gu et al. [9] have introduced a data-dependent SSM layer and developed Mamba, a universal language model backbone. Mamba exhibits superior performance compared to Transformers across various sizes when trained on large-scale real-world data while demonstrating linear scalability with sequence length. Based on this work, Huang et al. [13] solve the simple local 2D dependency relationship in the original Mamba model by dividing the image into different windows, effectively capturing local dependency relationships while maintaining Mamba’s [9] original global dependency ability. RS-Mamba [36] is proposed to process high-resolution remote sensing images. Vision-Mamba [37] combines bidirectional SSM for global visual context modeling for data dependency and position embedding for position-aware visual recognition.

State Space Models (SSMs) and various models based on Mamba are inspired by the continuous system, which maps the 1-D function or sequences $x(t)\!\in\!\mathbb{R}\rightarrow\!y(t)$ through a hidden state $h(t)$ $\in$ $\mathbb{R}^{N}$. Formally, SSMs employ the following ordinary differential equation (ODE) to model the input data:

where $\textbf{A}\!\in\!\mathbb{R}^{N\times N}$ represents the evolution matrix, $\textbf{B}\!\in\!\mathbb{R}^{N\times 1}$ and $\textbf{C}\!\in\!\mathbb{R}^{N\times 1}$ indicate the projection matrices. SSMs approximate this continuous ODE through discretization techniques. The S4 and Mamba are the discrete versions of the continuous system, which use the timescale parameter $\bm{\Delta}$ to convert continuous parameters $\mathbf{A}$ and $\mathbf{B}$ into discrete parameters $\overline{\mathbf{A}}$ and $\overline{\mathbf{B}}$. The commonly used method for transformation is zero-order hold (ZOH), which is defined as follows:

After the discretization of $\overline{\mathbf{A}}$ and $\overline{\mathbf{B}}$, the discretized version of Equation (1) using a step size $\bm{\Delta}$ can be rewritten as:

To enhance computational efficiency, the iterative process described in Equation (3) can be accelerated using parallel computation techniques, leveraging a global convolution operation:

where $\otimes$ is the convolution operation, $L$ is the length of the input sequence $\bm{x}$, and $\overline{\bm{K}}\!\in\!\mathbb{R}^{L}$ is the kernel of the SSM.

As shown in Fig. 1, CountMamba consists of three main components: Multi-directional State-Space Group (MSSG), Global-Local Adaptive Fusion (GLAF), Counter, and Normalizer. Given an input image $\mathbf{I}\!\in\!\mathbb{R}^{H\times W\times 3}$, we flatten it into a patch sequence and then project it to a feature sequence via linear transformation. Subsequently, we employ a Multi-directional State-Space Group (MSSG) to acquire the down-sampled deep features $\bm{F}_{H}$, $\bm{F}_{V}$, $\bm{F}_{D}$, and $\bm{F}_{A}$ from stacked Horizontal State-Space Blocks (HSSBs), Vertical State-Space Blocks (VSSBs), Diagonal State-Space Blocks (DSSBs) and Anti-diagonal State-Space Blocks (ASSBs). $\bm{F}_{H}$, $\bm{F}_{V}$, $\bm{F}_{D}$, and $\bm{F}_{A}$ are then fused adaptively in a sample-wise manner and refined by local features extracted by the CNN branch. Finally, the aggregated features undergo the Counter and Normalizer to predict the normalized counting map $\mathbf{C}_{n}\!\in\!\mathbb{R}^{H\times W}$, which is summed up to infer the overall number of plants $C$ in the image.

Multi-directional State-Space Group (MSSG) is designed to extract plant image features from various directions to adapt to different distributions. As shown in Fig. 1, MSSG is comprised of four parallel branches, consisting of stacked Horizontal State-Space Blocks (HSSBs), Vertical State-Space Blocks (VSSBs), Diagonal State-Space Blocks (DSSBs), and Anti-diagonal State-Space Blocks (ASSBs). HSSB, VSSB, DSSB, and ASSB share the same architecture except for the scan direction, possessing a Horizontal State-Space Module (HSSM), Vertical State-Space Module (VSSM), Diagonal State-Space Module (DSSM), Anti-diagonal State-Space Module (ASSM), respectively. The details of HSSM, VSSM, DSSM, and ASSM are illustrated in Fig. 2.

Inspired by the design of [36], a layer normalization is adopted to standardize the input data for balancing computational efficiency and capability, followed by a linear layer for transformation. Then, the depth-wise convolution operates on each channel separately to extract local features. The features are fed into State-Space Modules with corresponding directions to perform selective scanning in both the forward and backward directions, simulating an expert counting in a preferred scanning order. In addition, the output is linearly transformed and undergoes a gating operation with outputs of a linear transformation of the normalized features. The features undergo another linear layer and are added with inputs through a residual connection.

As plants in different images can be distributed in any direction, we design a sample-wise fusion mechanism to adaptively aggregate features from different directions, which can be defined as:

where $\bm{F}_{\text{global}}$ represents the adaptively aggregated feature embedding, the element-wise multiplication denoted by $\circ$, and $\bm{F}_{H},\bm{F}_{V},\bm{F}_{D},\bm{F}_{A}$ refer to the extracted featured from HSSM, VSSM, DSSM, and ASSM, respectively. The adaptive fusion weights $\alpha_{H}$, $\alpha_{V}$, $\alpha_{D}$, and $\alpha_{A}$ are defined as:

where $\mathbf{W}$ is a learnable linear transformation. Moreover, we add a lightweight CNN branch that contains stacked layers of convolution, batch normalization, non-linear activation, and max pooling to extract local features $\bm{F}_{\text{local}}$. The detailed architecture of the CNN branch is illustrated in Fig. 2. Finally, the aggregated features from multiple directions $\bm{F}_{\text{global}}$ are refined by $\bm{F}_{\text{local}}$ to obtain the final features $\bm{F}_{\text{fused}}$, defined as:

where $\beta$ is the hyper-parameter to control the weight of local information.

Counters are used to count plants in a specified area, while normalizers address the impact of duplicate counting of two overlapping images during the counting process [33]. After extracting the features of the image $\bm{F}_{\text{fused}}$, it needs to be sent to the counter for quantity prediction. In our experiment of the regression, a count map $C_{r}$ is used for the image transformation following [33], with each count representing a local region sized $r\!\times\!r$. During the image mapping process, the parameters of the local region size $r$ and the output stride $s$ are critical since they determine whether there is redundancy in the map. Throughout the experimental process, it is necessary to ensure that $r\!\geq\!s$. When $r\!>\!s$, every two adjacent local regions overlap with $\frac{r-s}{r}$, the count map becomes redundant in this situation. The overlap diminishes only when $r\!=\!s$. In our paper, the parameters of r and s are defined as 64 and 8, respectively. Thus, the resulting count map exhibits redundancy. The normalization step must be performed to eliminate redundancy and ensure that the final sum of normalized count maps accurately reflects the plant count of the image. The base input size of the network is $r\!\times\!r$, which is only tied to the network architecture. And owing to the architecture, the input size of the images can be arbitrary. For instance, the input image $\mathbf{I}\!\in\!\mathbb{R}^{H\times W\times 3}$, CountMamba defines a transformation $f$ as follows: $f(\mathbf{I}):\mathbb{R}^{H\times W\times 3}\!\rightarrow\!\mathbb{R}^{\frac{H}{s}\times\frac{W}{s}}$, where $H,W\gg r$.

As a redundant count map $\mathbf{C}_{r}\!\in\!\mathbb{R}^{\frac{H}{s}\times\frac{W}{s}}$ is emerged after the operation $f$, the normalized count map $\mathbf{C}_{n}\!\in\!\mathbb{R}^{H\times W}$ is designed. Subsequently, the image-level count $C$ can be calculated by aggregating $\mathbf{C}_{n}$, as follows:

where $\mathbf{C}_{n}(x,y)$ represents each element of $\mathbf{C}_{n}$ indexed by $x$ and $y$.

Following previous works [22, 33], we optimize our CountMamba with $L_{1}$ loss for counting values, which can be formulated as:

where $B$ is the batch size, $C_{i}$ and ${C}_{i}^{*}$ are the predicted number of plants and the ground truth in the $i$-th image, respectively.

MTC dataset is a maize tassels count dataset first introduced by Liu et al. [22]. The maize tassels in the MTC dataset exhibit significant variations in scale and shape. Their original resolution ranges from $3648\times 2736$, $4272\times 2848$, and $3456\times 2304$. For the experiment, 186 images are randomly selected and utilized as the training set, while the remaining 175 images are allocated to the test set. Every maize tassel in the dataset is manually labeled using a bounding box. These bounding-box annotations are then converted into dot annotations by calculating their central coordinates.

WED dataset is a wheat ear dataset first introduced by Madec et al. [24]. The images have a resolution of $6000\times 4000$ pixels, with the number of ears varying from 80 to 170 in each image. The dataset contains 236 images, with 165 allocated for training and 71 for testing. While bounding box annotations are provided in this dataset, only the center point of each box is utilized to ensure the uniform comparison of different methods for this experiment.

SHC dataset is a sorghum heads counting dataset first introduced by Guo et al. [24]. It consists of two subsets with 52 and 40 images, respectively. All the images are labeled with dotted annotations. Our experiments are evaluated on dataset 1, in which 26 images are randomly selected for training and the remaining for testing. We follow the same training configuration used in [33].

To evacuate the counting performance on these above datasets, we adopted the MAE, RMSE, rMAE, rRMSE, and $\text{R}^{2}$ as the basic metrics. These metrics are calculated as follows:

where $G_{i}$ is the real number of a plant in the ith image, $P_{i}$ is the prediction number of the plant in the ith image, $N$ is the number of total images, and $\bar{G}$ is the mean real number count.

All the images in the MTC dataset are resized to 512 $\times$ 512, and crop size is set to 256 $\times$ 256. As for the WED dataset, the down-sampling ratio is set to $\frac{1}{8}$, the crop size is set to 256 $\times$ 256. All the images in the SHC dataset are randomly cropped as 256 $\times$ 1024. The training process employs Adam optimizer with an initial learning rate of 1$\times$$10^{-4}$. The size of image patches is $2\times 2$, and the hyperparameter $\beta$ is set to 1. All experiments are carried out with the Pytorch framework and a single NVIDIA GeForce RTX 4090 GPU.

Tab. 1, 2, and 3 show the quantitative results between CountMamba and state-of-the-art plant counting methods on MTC, WED and SHC dataset, respectively. Thanks to the adaptation to the distribution of different directions, our proposed CountMamba achieves competitive performance on all three benchmark datasets. Notably, CountMamba achieves the best MAE of 4.6 and RMSE of 7.9 on the MTC dataset and the best RMSE of 18.6 on the SHC dataset.

Effects of different scan directions. To allow Mamba to count plants distributed in various directions, we build four SSMs inspired by [36], which uses scans in eight directions to generate scanned sequences. Here, we ablate different scan directions to study the effects, and the results are shown in 4. We are adding anti-diagonal, diagonal, vertical, and horizontal scanning in succession, which results in better results, demonstrating that counting in different order benefits the plant distribution in various directions.

Effects of expert numbers. We compare different numbers of counting experts in Tab. 5, including 1) One: one State-Space Block that scans in the horizontal, vertical, diagonal, and anti-diagonal directions simultaneously; 2) Two: one State-Space Block that scans in the horizontal and vertical directions and the other in the diagonal and anti-diagonal directions; 3) Four: using HSSB, VSSB, DSSB and ASSB. The results show that assigning feature extraction in the horizontal, vertical, diagonal, and anti-diagonal directions to four experts achieves the best counting results.

Effects of Global-Local Adaptive Fusion. We ablate two key designs in Global-Local Adaptive Fusion. The results, presented in 6, indicate that 1) Introducing a CNN branch to complement the global relationships with local information can effectively enhance the fine-grained local details. 2) Without using sample-wise adaptive fusion, i.e., directly using the average of HSSM, VSSM, DSSM, and ASSM for counting, can only obtain sub-optimal results. Using sample-wise adaptive fusion can flexibly adapt to plants with different spatial distributions.

Qualitative results of our proposed CountMamba on MTC, WED, and SHC datasets are shown in Fig. 3. We observe that: 1) CountMamba obtains strong responses on plant regions and weak responses on non-plant regions in the counting map and thus can infer accurate counting results. 2) CountMamba is robust to the direction of plant distribution and consistently effectively infers numbers from tens to hundreds.