| | --- |
| | license: mit |
| | language: |
| | - en |
| | metrics: |
| | - accuracy 0.9860 |
| | - roc_auc 0.9979 |
| | pipeline_tag: image-feature-extraction |
| | tags: |
| | - gnn |
| | - link-prediction |
| | - pytorch |
| | - pytorch-geometric |
| | - graph-neural-networks |
| | - biology |
| | - microscopy |
| | - cell-tracking |
| | - lineage-tracking |
| |
|
| | --- |
| | # DuMM Bacteria Tracker Model |
| | ## Overview |
| |
|
| | This repository contains the trained weights and documentation for the **DuMM Bacteria Tracker Model**, a **Graph Neural Network (GNN)** designed for **cell lineage link prediction** in time-lapse microscopy data. |
| |
|
| | The model uses a custom implementation of a **Parameter Decoupled Network (PDN)** variant within an **Edge-Propagation Message Passing Neural Network (EP-MPNN)** architecture, inspired by recent advancements in dynamic graph representation learning. |
| |
|
| | ## Model Architecture and Implementation Details |
| |
|
| | * **Framework:** PyTorch and PyTorch Geometric (PyG). |
| | * **Architecture:** `LineageLinkPredictionGNN` (custom `nn.Module`). |
| | * **Core Layers:** Utilizes custom `EP_MPNN_Block` which incorporates **Distance & Similarity (DS)** features and a **Jumping Knowledge (JK)** network for aggregating features across multiple layers. |
| | * **Task:** Binary classification (Link Prediction) on candidate edges between adjacent time frames (continuation or division links). |
| | * **Loss Function:** `BCEWithLogitsLoss`. |
| | * **Input Features:** Node features are scaled using a `StandardScalerTransform` and dynamically used to compute edge attributes (Absolute Difference + Cosine Similarity) via the `DS_block`. |
| |
|
| | ### Input Features |
| |
|
| | The model uses a **10-dimensional feature vector** for each node (cell) derived from image analysis. These features capture morphological and intensity properties of the bacterial cells across two channels (Phase Contrast and Fluorescence). |
| |
|
| | | Feature Name | Type | Description | |
| | | :--- | :--- | :--- | |
| | | `area` | Morphological | Area of the cell segment. | |
| | | `centroid_y` | Positional | Y-coordinate of the cell's centroid (critical for 1D growth systems). | |
| | | `axis_major_length` | Morphological | Length of the cell's major axis. | |
| | | `axis_minor_length` | Morphological | Length of the cell's minor axis. | |
| | | `intensity_mean/max/min_phase` | Intensity | Mean, max, and min pixel intensity in the **Phase Contrast** channel. | |
| | | `intensity_mean/max/min_fluor` | Intensity | Mean, max, and min pixel intensity in the **Fluorescence** channel. | |
| |
|
| |
|
| | The model was trained on microscopy images of the duplex mother machine developed by the Jun lab (https://jun.ucsd.edu/mother_machine.php) |
| | |
| | ### Data Preprocessing and Splitting |
| | |
| | #### Splitting Strategy |
| | To ensure the model generalizes to future, unseen data, a **time-based temporal split** was employed: |
| | 1. **Training Set:** First 60% of unique time frames (`sorted_time_frames[:train_split_idx]`). |
| | 2. **Validation Set:** Next 20% of unique time frames (`sorted_time_frames[train_split_idx:val_split_idx]`). |
| | 3. **Test Set:** Final 20% of unique time frames (`sorted_time_frames[val_split_idx:]`). |
| | |
| | #### Normalization |
| | * **Method:** Node features were normalized using **Standard Scaling (`sklearn.preprocessing.StandardScaler`)**. |
| | * **Fit:** The scaler was **fitted *only* on the training set features** (`all_train_node_features_df`). |
| | * **Application:** The fitted scaler was then applied to transform the features in the Training, Validation, and Test sets via the `StandardScalerTransform`. This avoids data leakage. |
| | |
| | #### Graph Creation (Candidate Generation) |
| | Candidate edges (links between cells in adjacent time frames) were generated based on a custom set of geometric and morphological heuristics: |
| | * **Distance Constraint:** Max distance between centroids is limited by `max_dist_link` (default 50.0). |
| | * **Area Ratio Constraints:** |
| | * **Continuation (1-to-1):** `min_area_ratio_continuation` (0.8) to `max_area_ratio_continuation` (1.2). |
| | * **Division (1-to-2):** `min_area_ratio_division` (1.8) to `max_area_ratio_division` (2.2). |
| | |
| | ### Training Protocol |
| | |
| | | Hyperparameter | Value | Description | |
| | | :--- | :--- | :--- | |
| | | **GNN Layers (`num_blocks`)** | 2 | Number of sequential EP-MPNN blocks. | |
| | | **Hidden Channels** | 128 | Dimension for node and edge embeddings. | |
| | | **Optimizer** | Adam | Standard optimization algorithm. | |
| | | **Learning Rate** | 0.001 | Base learning rate. | |
| | | **Weight Decay** | 0.0005 | L2 regularization applied to prevent overfitting. | |
| | | **Batch Size** | 32 | Number of graphs processed per iteration. | |
| | | **Evaluation Metric** | Validation Accuracy (`val_acc`) | Used for saving the `best_link_prediction_model.pt`. | |
| | | **Early Stopping** | Yes | Monitors Validation Loss (`val_loss`) with a **patience of 10 epochs**. | |
| | | **Max Epochs** | 500 | Maximum number of training epochs. | |
