Edited by: Jiajie Peng, Northwestern Polytechnical University, China
Reviewed by: Xishuang Dong, Prairie View A&M University, United States; Naifeng Wen, Dalian Nationalities University, China; Zhong Li, Institute of Disaster Prevention, China
This article was submitted to Bioinformatics and Computational Biology, a section of the journal Frontiers in Bioengineering and Biotechnology
This is an open-access article distributed under the terms of the Creative Commons Attribution License (CC BY). The use, distribution or reproduction in other forums is permitted, provided the original author(s) and the copyright owner(s) are credited and that the original publication in this journal is cited, in accordance with accepted academic practice. No use, distribution or reproduction is permitted which does not comply with these terms.
Accurate target detection and association are vital for the development of reliable target tracking, especially for cell tracking based on microscopy images due to the similarity of cells. We propose a deep reinforcement learning method to associate the detected targets between frames. According to the dynamic model of each target, the cost matrix is produced by conjointly considering various features of targets and then used as the input of a neural network. The proposed neural network is trained using reinforcement learning to predict a distribution over the association solution. Furthermore, we design a residual convolutional neural network that results in more efficient learning. We validate our method on two applications: the multiple target tracking simulation and the ISBI cell tracking. The results demonstrate that our approach based on reinforcement learning techniques could effectively track targets following different motion patterns and show competitive results.
Tracking individual cells in a group is the fundamental of many biomedical analysis tasks, including understanding how genotypes are related to phenotypes, tracking the early development of organs and meristems, and potentially tracking the development of cancerous tumors (Cheng et al.,
There are many procedures and methods for tracking objects at the microscopic level. Tracking-by-detection methods are widely used in multi-target tracking, in which detection and association are two primary issues. Extensive research efforts have focused on detection, especially in cell-tracking applications. For target association between frames, the naïve nearest-neighbor method is commonly adopted but provides unsatisfactory association accuracy. Target association is a combinatorial optimization problem, which is widely studied in computer science and mathematics and many such problems are NP-hard. In general, the linear assignment problem is to find the optimal assignment that maximizes or minimizes the sum of the costs in a cost matrix. Classic algorithms for the linear assignment include the Hungarian method (Kuhn,
Recently, some data-driven methods have been proposed to solve combinatorial problems. Vinyals et al. (
Substantial progress in artificial intelligence has been made in supervised learning, where systems are trained on vast amounts of labeled data (Peng et al.,
This work is motivated by several recent proposed DRL methods for NP-hard problems. We propose a DRL approach to automatically search for assignment solutions for a given cost matrix. Specifically, we first modeled the association of cells between frame as an linear assignment problem and formulated the assignment problem with the one-to-one constraint as a DRL problem. Then, with the objective of minimizing the sum of the assignment costs, we used DRL to obtain the optimal assignment solution. To convert the cost matrix into a finite action space, we employ the residual learning and convolutional neural network (CNN) to extract features from a set of training samples and use the pointing mechanism (Bello et al.,
Our contributions are the following:(1) A simple framework for cell detection and association based on the idea of (2) We introduce a formulation that translates the decision making in the linear assignment problem algorithm into an RL problem. (3) We propose a novel neural network architecture that end-to-end maps the inputs to the decision outputs.
The organization of this paper is as follows. Related work is introduced in section 2. The framework of the proposed method and training details are presented in section 3. In section 4, some experiments are conducted to evaluate the performance of our proposed method. The conclusion is given in section 5.
A large variety of cell tracking methods have been described in the existing literature. These cell tracking methods can be broadly grouped into two categories: (i) tracking by model evolution and (ii) tracking by detection.
In tracking by model evolution methods, cell segmentation and tracking are solved simultaneously in each frame of a cell video. Typically, these methods are driven by data in some feature space and make a regularity assumption on the smoothness of the curve. In this framework, cells are represented by parametric or implicit active contour models. Parametric models utilize the explicit representations of cell boundaries such as Gaussian Mixture Models (GMM) (Amat et al.,
Existing cell tracking methods generally adopt the tracking by detection strategy. The tracking by detection method typically consists of two stages: the cell detection stage and cell association stage. In the first stage, the cells are detected by image segmentation methods. Subsequently, in the second stage, detected cells are associated with neighboring frames in real-time or all frames offline. Cell detection can be achieved by classic image segmentation algorithms based on intensity features, gradient features, or texture features (Chenouard et al.,
In this section we present a tracking by detection approach to construct the cell trajectories from a time-series microscopy image sequence. The framework consists of two modules: cell detection and cell association. The U-Net segmentation method is employed to detect all the cells in each frame, and then we adopt the traditional single hypothesis tracking method with Kalman filter and frame-by-frame data association to produce the cell trajectories.
Ronneberger et al. (
In this work, we assume that each cell can be modeled as a discrete-time Markov process:
where
A Kalman filter can be adopted to use those cell detection results to predict the state of cells, which can then be used to formulate the cell association between frames as a linear assignment problem.
To solve the target association problem by DRL, we present our solution architecture in three parts: (1) Problem Formulation. We formulate the procedure for selecting an assignment solution as an RL problem to associate target states and measurements. (2) Neural network architecture. An end-to-end architecture that maps from the state space to the action space is designed. (3) Training algorithm. We present the RL algorithm used for the policy search.
Assume that the cell trajectories can be denoted as a set
The values of the cost matrix
where
Illustration of the proposed association matrix.
The standard RL formulation starts with an MDP: at time step
To formulate the procedure of selecting assignment solution algorithms into an MDP, we specify below the state space
The input of our residual CNN (ResCNN) is a cost matrix
The sequence-to-sequence models the linear assignment problem (Milan et al.,
The architecture of the proposed ResCNN network.
Our proposed neural network is similar to the image denoising network introduced in Zhang et al. (
In summary, the main feature of our ResCNN is the adoption of residual learning to learn
In the following, we will give some important details about our network design and training.
Batch normalization is a standard technique that is widely used in image classification CNN models. Training a deep neural network model is often difficult not only because of the gradient vanishing/exploding problem but also because the distribution of data changes between layers, which is called the “internal covariate shift" phenomenon. Batch normalization is a technique that can relieve this phenomenon by introducing several simple operations to the input data. The goal of the normalization step for batch normalization is to transform the layer input
where E[
Batch normalization and residual learning are two important algorithms for designing a neural network architecture. Residual learning and batch normalization can benefit from each other (Zhang et al.,
In the linear assignment problem, the input and output need to be consistent. However, due to the characteristics of convolution, the neural network is prone to producing boundary artifacts without proper handling. There are two common ways to solve this problem: symmetrical padding and zero padding. In our work, we select zero padding to maintain a consistent matrix size.
Unlike ordinary visual tasks, for the linear assignment problem, one major characteristic is that one detected cell can only be assigned to one predicted cell. The neural network output should satisfy one-to-one constraints. Let
where
In this paper, we utilize the RL to train the neural network. The input of the network can be denoted as
In our work,
where
For a cost matrix, the baseline value
Algorithm 1 gives the pseudo-code of the training procedure of the neural network.
Training Procedure
1: Training set |
2: Initialize the neural net params θ. |
3: Initialize baseline value. |
4: |
5: Select a batch of samples |
6: Sample solution π |
7: Let |
8: Update θ = |
9: Update baseline |
10: |
11: return neural net parameters θ. |
To evaluate the performance of our proposed method, we consider two applications of the linear assignment problem: maximum weight matching (MWM) and data association for multi-target tracking. We first compare our method with the state-of-the-art DRL method for maximum weight matching. Then, we test our method on a multi-target tracking scenario. Finally, we evaluate our proposed method on three cell microscopy datasets, Fluo-N2DH-GOWT1, PhC-C2DH-U373, and Fluo-N2DH-SIM+ from the ISBI 2015 Cell Tracking Challenge (Maška et al.,
In all experiments, we used 500, 000 training samples for the data association. To produce the training samples, we randomly sample
Define a weighted bipartite graph
We trained our method on MWM with
Median optimality ratios on the MWM test set.
AC+Matching | 0.935 | 0.897 | 0.725 |
SPG+Matching | 0.904 | 0.895 | 0.889 |
Ours | 0.977 | 0.968 | 0.965 |
One major application of linear assignment is data association for multi-target tracking. Therefore, we set up a simulated multi-target tracking scenario to evaluate the performance of the proposed method similar to Milan et al. (
Comparison of the track maintenance performance of different algorithms:
We replace the data association part of JPDA with our method and call it JPDA-RL. The input matrix
We employ two metrics to evaluate the tracking results: the Optimal Sub-pattern Assignment metric for track (OSPA-T) and Number of Identity Switch (IDSW). The OSPA-T distance (Ristic et al.,
where
To compute the OSPA-T distance for the estimated tracks and true tracks, two parameters, the cardinality penalty
In
Average OSPA-T distance and IDSW for different methods over 100 random runs.
JPDA | 0.19(0.05) | 0.90(0.88) | 0.34(0.11) | 0.70(0.82) | 0.41(0.12) | 0.40(0.70) |
JPDA10 | 0.23(0.10) | 0.70(0.67) | 0.37(0.11) | 0.90(0.88) | 0.43(0.09) | 1.10(0.99) |
JPDA-HA | 0.28(0.06) | 0.60(0.84) | 0.37(0.10) | 0.70(0.95) | 0.46(0.14) | 1.30(0.82) |
JPDA-RL | 0.28(0.06) | 0.60(0.84) | 0.36(0.08) | 0.60(0.70) | 0.45(0.13) | 1.10(0.99) |
LSTM | 0.11(0.01) | 1.07(0.84) | 0.21(0.01) | 1.00(0.74) | 0.37(0.11) | 0.60(0.89) |
The segmentation task by U-Net and data association by DRL are conducted on AMD Ryzen 9 3900X 12 core processors with a GeForce GTX 2060 graphics card. For comparison, segmentation (SEG), tracking (TRA) accuracy measures and overall performance (OP) are adopted to evaluate the tracking performance. For TRA, Acyclic Oriented Graph Matching (AOGM) is used to count the changes needed to transform the cell tracking family tree into the ground-truth graph. OP is defined as the mean of TRA and SEG.
The results of this work are compared against the best performing available methods for each dataset. For the
TRA, SEG and OPT performance for our method, CPN, KTH (Magnusson and Jaldén,
Fluo-N2DH-GOWT1-01 | CPN | 0.9864 | 0.8506 | 0.9185 |
BLOB | 0.9733 | 0.7415 | 0.8574 | |
KTH | 0.9462 | 0.6849 | 0.8155 | |
Ours | ||||
Fluo-N2DH-GOWT1-02 | CPN | 0.8725 | 0.9222 | |
BLOB | 0.9628 | 0.9046 | 0.9337 | |
KTH | 0.9452 | 0.8942 | 0.9197 | |
Ours | 0.9575 | |||
PhC-C2DH-U373-01 | CPN | 0.9594 | 0.7336 | 0.8456 |
U-Net | 0.9869 | |||
GC-ME | 0.9779 | 0.8748 | 0.9264 | |
Ours | 0.8527 | 0.9223 | ||
PhC-C2DH-U373-02 | CPN | 0.9346 | 0.7376 | 0.8361 |
U-Net | 0.8925 | |||
GC-ME | 0.9040 | 0.7567 | 0.8304 | |
Ours | 0.9318 | 0.7735 | 0.8527 | |
Fluo-N2DH-SIM+-01 | U-Net-S | |||
Ours | 0.9841 | 0.8854 | 0.9348 | |
Fluo-N2DH-SIM+-02 | U-Net-S | 0.9597 | 0.7381 | 0.8489 |
Ours |
For the
For
In this paper, we presented a solution to the problem of data association in cell tracking using the deep reinforcement learning. We formulated the data association problem into a linear assignment problem and then proposed a deep reinforcement learning framework which utilizes a residual CNN neural network. In simulation results, we compare the proposed method with other state-of-the-art approaches on various cell tracking datasets, and the results show that the proposed method achieves better comprehensive performance. Thus, our method likely has applications in the field of biomedical engineering. There are also some limitations of our tracking method that leave room for improvement. In future research, we plan to improve the data association method to deal with one-to-many and many-to-one association problems.
All datasets generated for this study are included in the article/supplementary material.
LZ and JW substantially contributed to the conception and design of the study. XS analyzed and interpreted the data. LZ, JW, and JZ drafted the article.
The authors declare that the research was conducted in the absence of any commercial or financial relationships that could be construed as a potential conflict of interest.
We thank all of the reviewers for their valuable comments and suggestions.