Add MEDIAR models for cell segmentation

annolid/gui/app.py

@@ -315,6 +315,14 @@ def __init__(self,
             self.tr("Open video")
         )
 
+        segment_cells = action(
+            self.tr("&Segment Cells"),
+            self._segment_cells,
+            None,
+            "Segment Cells",
+            self.tr("Segment Cells")
+        )
+
         advance_params = action(
             self.tr("&Advanced Parameters"),
             self.set_advanced_params,
@@ -328,6 +336,12 @@ def __init__(self,
             )
         ))
 
+        segment_cells.setIcon(QtGui.QIcon(
+            str(
+                self.here / "icons/cell_seg.png"
+            )
+        ))
+
         open_audio = action(
             self.tr("&Open Audio"),
             self.openAudio,
@@ -488,6 +502,7 @@ def __init__(self,
         _action_tools.append(quality_control)
         _action_tools.append(colab)
         _action_tools.append(visualization)
+        _action_tools.append(segment_cells)
 
         self.actions.tool = tuple(_action_tools)
         self.tools.clear()
@@ -502,6 +517,7 @@ def __init__(self,
         utils.addActions(self.menus.file, (models,))
         utils.addActions(self.menus.file, (tracks,))
         utils.addActions(self.menus.file, (quality_control,))
+        utils.addActions(self.menus.file, (segment_cells,))
         utils.addActions(self.menus.file, (downsample_video,))
         utils.addActions(self.menus.file, (convert_sleap,))
         utils.addActions(self.menus.file, (advance_params,))
@@ -991,6 +1007,46 @@ def _select_sam_model_name(self):
 
         return model_name
 
+    def _segment_cells(self):
+        if self.filename or len(self.imageList) > 0:
+            from annolid.segmentation.MEDIAR.predict_ensemble import MEDIARPredictor
+            if self.annotation_dir is not None:
+                out_dir_path = self.annotation_dir + '_masks'
+            elif self.filename:
+                out_dir_path = str(Path(self.filename).with_suffix(''))
+            if not os.path.exists(out_dir_path):
+                os.makedirs(out_dir_path, exist_ok=True)
+
+            if self.filename is not None and self.annotation_dir is None:
+                self.annotation_dir = out_dir_path
+                target_link = os.path.join(
+                    out_dir_path, os.path.basename(self.filename))
+                if not os.path.islink(target_link):
+                    os.symlink(self.filename, target_link)
+                else:
+                    logger.info(f"The symlink {target_link} alreay exists.")
+            mediar_predictor = MEDIARPredictor(input_path=self.annotation_dir,
+                                               output_path=out_dir_path)
+            self.worker = FlexibleWorker(
+                mediar_predictor.conduct_prediction)
+            self.thread = QtCore.QThread()
+            self.worker.moveToThread(self.thread)
+            self.worker.start.connect(self.worker.run)
+            self.worker.finished.connect(self.thread.quit)
+            self.worker.finished.connect(self.worker.deleteLater)
+            self.thread.finished.connect(self.thread.deleteLater)
+            self.worker.finished.connect(lambda:
+                                         QtWidgets.QMessageBox.about(self,
+                                                                     "Cell counting results are ready",
+                                                                     f"Please review your results."))
+            self.worker.return_value.connect(
+                lambda shape_list: self.loadShapes(shape_list))
+
+            self.thread.start()
+            self.worker.start.emit()
+            # shape_list = mediar_predictor.conduct_prediction()
+            # self.loadShapes(shape_list)
+
     def stop_prediction(self):
         # Emit the stop signal to signal the prediction thread to stop
         self.pred_worker.stop()
@@ -1009,7 +1065,7 @@ def predict_from_next_frame(self,
         if self.pred_worker and self.stop_prediction_flag:
             # If prediction is running, stop the prediction
             self.stop_prediction()
-        elif len(self.canvas.shapes) <= 0:
+        elif len(self.canvas.shapes) <= 0 and self.video_file is not None:
             QtWidgets.QMessageBox.about(self,
                                         "No Shapes or Labeled Frames",
                                         f"Please label this frame")

annolid/segmentation/MEDIAR/.gitignore

@@ -0,0 +1,18 @@
+config/
+*.log
+*.ipynb
+*.ipynb_checkpoints/
+__pycache__/
+results/
+weights/
+wandb/
+data/
+submissions/
+/.vscode
+*.npy
+*.pth
+*.sh
+*.json
+*.out
+*.zip
+*.tiff

annolid/segmentation/MEDIAR/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 opcrisis
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

annolid/segmentation/MEDIAR/README.md

@@ -0,0 +1,189 @@
+
+# **MEDIAR: Harmony of Data-Centric and Model-Centric for Multi-Modality Microscopy**
+![1-2](https://user-images.githubusercontent.com/12638561/207771867-0b1414f2-cf48-4747-9cda-3304e6d86bfd.png)
+
+
+This repository provides an official implementation of [MEDIAR: MEDIAR: Harmony of Data-Centric and Model-Centric for Multi-Modality Microscopy](https://arxiv.org/abs/2212.03465), which achieved the ***"1st winner"*** in the [NeurIPS-2022 Cell Segmentation Challenge](https://neurips22-cellseg.grand-challenge.org/).
+
+To access and try mediar directly, please see links below.
+- <a href="https://colab.research.google.com/drive/1iFnGu6A_p-5s_eATjNtfjb9-MR5L3pLB?usp=sharing"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
+- [Huggingface Space](https://huggingface.co/spaces/ghlee94/MEDIAR?logs=build)
+- [Napari Plugin](https://github.com/joonkeekim/mediar-napari)
+- [Docker Image](https://hub.docker.com/repository/docker/joonkeekim/mediar/general)
+
+# 1. MEDIAR Overview
+<img src="./image/mediar_framework.png" width="1200"/>
+
+MEIDAR is a framework for efficient cell instance segmentation of multi-modality microscopy images. The above figure illustrates an overview of our approach. MEDIAR harmonizes data-centric and model-centric approaches as the learning and inference strategies, achieving a **0.9067** Mean F1-score on the validation datasets. We provide a brief description of methods that combined in the MEDIAR. Please refer to our paper for more information.
+# 2. Methods
+
+## **Data-Centric**
+- **Cell Aware Augmentation** : We apply two novel cell-aware augmentations. *cell-wisely intensity is randomization* (Cell Intensity Diversification) and *cell-wise boundary pixels exclusion* in the label. The boundary exclusion is adopted only in the pre-training phase.
+
+- **Two-phase Pretraining and Fine-tuning** : To extract knowledge from large public datasets, we first pretrained our model on public sets, then fine-tune.
+    - Pretraining : We use 7,2412 labeled images from four public datasets for pretraining: OmniPose, CellPose, LiveCell and DataScienceBowl-2018. MEDIAR takes two different phases for the pretraining. the MEDIAR-Former model with encoder parameters initialized from ImageNet-1k pretraining.
+
+    - Fine-tuning : We use two different model for ensemble. First model is fine-tuned 200 epochs using target datasets. Second model is fine-tuned 25 epochs using both target and public datsets.
+
+- **Modality Discovery & Amplified Sampling** : To balance towards the latent modalities in the datasets, we conduct K-means clustering and discover 40 modalities. In the training phase, we over-sample the minor cluster samples.
+
+- **Cell Memory Replay** : We concatenate the data from the public dataset with a small portion to the batch and train with boundary-excluded labels.
+
+## **Model-Centric**
+- **MEDIAR-Former Architecture** : MEDIAR-Former follows the design paradigm of U-Net, but use SegFormer and MA-Net for the encoder and decoder. The two heads of MEDIAR-Former predicts cell probability and gradieng flow. 
+
+<img src="./image/mediar_model.PNG" width="1200"/>
+
+- **Gradient Flow Tracking** : We utilize gradient flow tracking proposed by [CellPose](https://github.com/MouseLand/cellpose).
+
+- **Ensemble with Stochastic TTA**: During the inference, the MEIDAR conduct prediction as sliding-window manner with importance map generated by the gaussian filter. We use two fine-tuned models from phase1 and phase2 pretraining, and ensemble their outputs by summation. For each outputs, test-time augmentation is used.
+# 3. Experiments
+
+### **Dataset**
+- Official Dataset
+  - We are provided the target dataset from [Weakly Supervised Cell Segmentation in Multi-modality High-Resolution Microscopy Images](https://neurips22-cellseg.grand-challenge.org/). It consists of 1,000 labeled images, 1,712 unlabeled images and 13 unlabeled whole slide image from various microscopy types, tissue types, and staining types. Validation set is given with 101 images including 1 whole slide image.
+
+- Public Dataset
+  - [OmniPose](http://www.cellpose.org/dataset_omnipose) : contains mixtures of 14 bacterial species. We only use 611 bacterial cell microscopy images and discard 118 worm images.
+  - [CellPose](https://www.cellpose.org/dataset) : includes Cytoplasm, cellular microscopy, fluorescent cells images. We used 551 images by discarding 58 non-microscopy images. We convert all images as gray-scale.
+  - [LiveCell](https://github.com/sartorius-research/LIVECell) : is a large-scale dataset with 5,239 images containing 1,686,352 individual cells annotated by trained crowdsources from 8 distinct cell types.
+  - [DataScienceBowl 2018](https://www.kaggle.com/competitions/sartorius-cell-instance-segmentation/overview) : 841 images contain 37,333 cells from 22 cell types, 15 image resolutions, and five visually similar groups.
+
+### **Testing steps**
+- **Ensemble Prediction with TTA** : MEDIAR uses sliding-window inference with the overlap size between the adjacent patches as 0.6 and gaussian importance map. To predict the different views on the image, MEDIAR uses Test-Time Augmentation (TTA) for the model prediction and ensemble two models described in **Two-phase Pretraining and Fine-tuning**.
+
+- **Inference time** : MEDIAR conducts most images in less than 1sec and it depends on the image size and the number of cells, even with ensemble prediction with TTA. Detailed evaluation-time results are in the paper.  
+
+### **Preprocessing & Augmentations**
+| Strategy  |      Type      |  Probability |
+|----------|:-------------|------|
+| `Clip` |  Pre-processing | . |
+| `Normalization` |  Pre-processing | . |
+| `Scale Intensity` |  Pre-processing | . |
+| `Zoom` |  Spatial Augmentation | 0.5 |
+| `Spatial Crop` |  Spatial Augmentation | 1.0 |
+| `Axis Flip` |  Spatial Augmentation | 0.5 |
+| `Rotation` |  Spatial Augmentation | 0.5 |
+| `Cell-Aware Intensity` |  Intensity Augmentation | 0.25 |
+| `Gaussian Noise` |  Intensity Augmentation | 0.25 |
+| `Contrast Adjustment` |  Intensity Augmentation | 0.25 |
+| `Gaussian Smoothing` |  Intensity Augmentation | 0.25 |
+| `Histogram Shift` |  Intensity Augmentation | 0.25 |
+| `Gaussian Sharpening` |  Intensity Augmentation | 0.25 |
+| `Boundary Exclusion` | Others | . |
+
+
+| Learning Setups                                             | Pretraining                                    | Fine-tuning                                    |
+|----------------------------------------------------------------------|---------------------------------------------------------|---------------------------------------------------------|
+| Initialization (Encoder)                                             | Imagenet-1k pretrained                                  | from Pretraining                                              |
+| Initialization (Decoder, Head)                                       |            He normal initialization                     | from Pretraining|
+| Batch size                                                           | 9                                                       | 9                                                       |
+| Total epochs                                                         | 80 (60)                                                     | 200 (25)                                                     |
+| Optimizer                                                            | AdamW                                     | AdamW                                     |
+| Initial learning rate (lr)                                           | 5e-5                                                    | 2e-5                                                    |
+| Lr decay schedule                                                    | Cosine scheduler (100 interval) | Cosine scheduler (100 interval) |
+| Loss function                                                        | MSE, BCE                              | MSE, BCE                               |
+
+# 4. Results
+### **Validation Dataset**
+- Quantitative Evaluation
+    - Our MEDIAR achieved **0.9067** validation mean F1-score.
+- Qualitative Evaluation
+<img src="./image/mediar_results.png" width="1200"/>
+
+- Failure Cases
+<img src="./image/failure_cases.png" width="1200"/>
+
+### **Test Dataset**
+![F1_osilab](https://user-images.githubusercontent.com/12638561/207772559-2185b79c-8288-4556-a3b4-9bd1d359fceb.png)
+![RunningTime_osilab](https://user-images.githubusercontent.com/12638561/207772555-c3b29071-6e03-4985-837a-da7b3dd3b65d.png)
+
+
+# 5. Reproducing
+
+### **Our Environment**
+| Computing Infrastructure| |
+|-------------------------|----------------------------------------------------------------------|
+| System                  | Ubuntu 18.04.5 LTS                                                   |
+| CPU                     | AMD EPYC 7543 32-Core Processor [email protected]                          |
+| RAM                     | 500GB; 3.125MT/s                                                     |
+| GPU (number and type)   | NVIDIA A5000 (24GB) 2ea                                              |
+| CUDA version            | 11.7                                                                 |
+| Programming language    | Python 3.9                                                           |
+| Deep learning framework | Pytorch (v1.12, with torchvision v0.13.1)             |
+| Code dependencies       | MONAI (v0.9.0), Segmentation Models (v0.3.0) |
+| Specific dependencies   | None                                                                 |
+
+To install requirements:
+
+```
+pip install -r requirements.txt
+wandb off
+```
+
+## Dataset
+- The datasets directories under the root should the following structure:
+
+```
+  Root
+  ├── Datasets
+  │   ├── images (images can have various extensions: .tif, .tiff, .png, .bmp ...)
+  │   │    ├── cell_00001.png
+  │   │    ├── cell_00002.tif
+  │   │    ├── cell_00003.xxx
+  │   │    ├── ...  
+  │   └── labels (labels must have .tiff extension.)
+  │   │    ├── cell_00001_label.tiff 
+  │   │    ├── cell_00002.label.tiff
+  │   │    ├── cell_00003.label.tiff
+  │   │    ├── ...
+  └── ...
+```
+
+Before execute the codes, run the follwing code to generate path mappting json file:
+
+```python
+python ./generate_mapping.py --root=<path_to_data>
+```
+
+## Training
+
+To train the model(s) in the paper, run the following command:
+
+```python
+python ./main.py --config_path=<path_to_config>
+```
+Configuration files are in `./config/*`. We provide the pretraining, fine-tuning, and prediction configs. You can refer to the configuration options in the `./config/mediar_example.json`. We also implemented the official challenge baseline code in our framework. You can run the baseline code by running the `./config/baseline.json`.
+
+## Inference
+
+To conduct prediction on the testing cases, run the following command:
+
+```python
+python predict.py --config_path=<path_to_config>
+```
+
+## Evaluation
+If you have the labels run the following command for evaluation:
+
+```python
+python ./evaluate.py --pred_path=<path_to_prediciton_results> --gt_path=<path_to_ground_truth_labels>
+```
+
+The configuration files for `predict.py` is slightly different. Please refer to the config files in `./config/step3_prediction/*`.
+## Trained Models
+
+You can download MEDIAR pretrained and finetuned models here:
+
+- [Google Drive Link](https://drive.google.com/drive/folders/1RgMxHIT7WsKNjir3wXSl7BrzlpS05S18?usp=share_link). 
+
+## Citation of this Work
+```
+@article{lee2022mediar,
+  title={Mediar: Harmony of data-centric and model-centric for multi-modality microscopy},
+  author={Lee, Gihun and Kim, SangMook and Kim, Joonkee and Yun, Se-Young},
+  journal={arXiv preprint arXiv:2212.03465},
+  year={2022}
+}
+```
+

annolid/segmentation/MEDIAR/SetupDict.py

@@ -0,0 +1,39 @@
+import torch.optim as optim
+import torch.optim.lr_scheduler as lr_scheduler
+import monai
+
+from annolid.segmentation.MEDIAR import core
+from annolid.segmentation.MEDIAR.train_tools import models
+from annolid.segmentation.MEDIAR.train_tools.models import *
+
+__all__ = ["TRAINER", "OPTIMIZER", "SCHEDULER"]
+
+TRAINER = {
+    "baseline": core.Baseline.Trainer,
+    "mediar": core.MEDIAR.Trainer,
+}
+
+PREDICTOR = {
+    "baseline": core.Baseline.Predictor,
+    "mediar": core.MEDIAR.Predictor,
+    "ensemble_mediar": core.MEDIAR.EnsemblePredictor,
+}
+
+MODELS = {
+    "unet": monai.networks.nets.UNet,
+    "unetr": monai.networks.nets.unetr.UNETR,
+    "swinunetr": monai.networks.nets.SwinUNETR,
+    "mediar-former": models.MEDIARFormer,
+}
+
+OPTIMIZER = {
+    "sgd": optim.SGD,
+    "adam": optim.Adam,
+    "adamw": optim.AdamW,
+}
+
+SCHEDULER = {
+    "step": lr_scheduler.StepLR,
+    "multistep": lr_scheduler.MultiStepLR,
+    "cosine": lr_scheduler.CosineAnnealingLR,
+}