spacr 0.2.81__py3-none-any.whl → 0.3.1__py3-none-any.whl

spacr/mediar.py

@@ -0,0 +1,364 @@
+import os, sys, gdown, cv2, torch
+import numpy as np
+import matplotlib.pyplot as plt
+from monai.inferers import sliding_window_inference
+import skimage.io as io
+
+# Path to the MEDIAR directory
+mediar_path = os.path.join(os.path.dirname(__file__), 'resources', 'MEDIAR')
+
+print('mediar path', mediar_path)
+
+# Temporarily create __init__.py to make MEDIAR a package
+init_file = os.path.join(mediar_path, '__init__.py')
+if not os.path.exists(init_file):
+    with open(init_file, 'w'):  # Create the __init__.py file
+        pass
+
+# Add MEDIAR to sys.path
+sys.path.insert(0, mediar_path)
+
+try:
+    # Now import the dependencies from MEDIAR
+    from core.MEDIAR import Predictor, EnsemblePredictor
+    from train_tools.models import MEDIARFormer
+finally:
+    # Remove the temporary __init__.py file after the import
+    if os.path.exists(init_file):
+        os.remove(init_file)  # Remove the __init__.py file
+
+def display_imgs_in_list(lists_of_imgs, cmaps=None):
+    """
+    Displays images from multiple lists side by side. 
+    Each row will display one image from each list (lists_of_imgs[i][j] is the j-th image in the i-th list).
+    
+    :param lists_of_imgs: A list of lists, where each inner list contains images.
+    :param cmaps: List of colormaps to use for each list (optional). If not provided, defaults to 'gray' for all lists.
+    """
+    num_lists = len(lists_of_imgs)
+    num_images = len(lists_of_imgs[0])
+    
+    # Ensure that all lists have the same number of images
+    for img_list in lists_of_imgs:
+        assert len(img_list) == num_images, "All inner lists must have the same number of images"
+
+    # Use 'gray' as the default colormap if cmaps are not provided
+    if cmaps is None:
+        cmaps = ['gray'] * num_lists
+    else:
+        assert len(cmaps) == num_lists, "The number of colormaps must match the number of lists"
+
+    plt.figure(figsize=(15, 5 * num_images))
+    
+    for j in range(num_images):
+        for i, img_list in enumerate(lists_of_imgs):
+            img = img_list[j]
+            plt.subplot(num_images, num_lists, j * num_lists + i + 1)
+            
+            if len(img.shape) == 2:  # Grayscale image
+                plt.imshow(img, cmap=cmaps[i])
+            elif len(img.shape) == 3 and img.shape[0] == 3:  # 3-channel image (C, H, W)
+                plt.imshow(img.transpose(1, 2, 0))  # Change shape to (H, W, C) for displaying
+            else:
+                plt.imshow(img)
+            
+            plt.axis('off')
+            plt.title(f'Image {j+1} from list {i+1}')
+
+    plt.tight_layout()
+    plt.show()
+
+def get_weights(finetuned_weights=False):
+    if finetuned_weights:
+        model_path1 = os.path.join(os.path.dirname(__file__), 'resources', 'MEDIAR_weights', 'from_phase1.pth')
+        if not os.path.exists(model_path1):
+            print("Downloading finetuned model 1...")
+            gdown.download('https://drive.google.com/uc?id=1JJ2-QKTCk-G7sp5ddkqcifMxgnyOrXjx', model_path1, quiet=False)
+    else:
+        model_path1 = os.path.join(os.path.dirname(__file__), 'resources', 'MEDIAR_weights', 'phase1.pth')
+        if not os.path.exists(model_path1):
+            print("Downloading model 1...")
+            gdown.download('https://drive.google.com/uc?id=1v5tYYJDqiwTn_mV0KyX5UEonlViSNx4i', model_path1, quiet=False)
+            
+    if finetuned_weights:
+        model_path2 = os.path.join(os.path.dirname(__file__), 'resources', 'MEDIAR_weights', 'from_phase2.pth')
+        if not os.path.exists(model_path2):
+            print("Downloading finetuned model 2...")
+            gdown.download('https://drive.google.com/uc?id=168MtudjTMLoq9YGTyoD2Rjl_d3Gy6c_L', model_path2, quiet=False)
+    else:
+        model_path2 = os.path.join(os.path.dirname(__file__), 'resources', 'MEDIAR_weights', 'phase2.pth')
+        if not os.path.exists(model_path2):
+            print("Downloading model 2...")
+            gdown.download('https://drive.google.com/uc?id=1NHDaYvsYz3G0OCqzegT-bkNcly2clPGR', model_path2, quiet=False)
+    
+    return model_path1, model_path2
+
+def normalize_image(image, lower_percentile=0.0, upper_percentile=99.5):
+    """
+    Normalize an image based on the 0.0 and 99.5 percentiles.
+    
+    :param image: Input image (numpy array).
+    :param lower_percentile: Lower percentile (default is 0.0).
+    :param upper_percentile: Upper percentile (default is 99.5).
+    :return: Normalized image (numpy array).
+    """
+    lower_bound = np.percentile(image, lower_percentile)
+    upper_bound = np.percentile(image, upper_percentile)
+    
+    # Clip image values to the calculated percentiles
+    image = np.clip(image, lower_bound, upper_bound)
+    
+    # Normalize to [0, 1]
+    image = (image - lower_bound) / (upper_bound - lower_bound + 1e-5)  # Add small epsilon to avoid division by zero
+    
+    return image
+
+class MEDIARPredictor:
+    def __init__(self, input_path=None, output_path=None, device=None, model="ensemble", roi_size=512, overlap=0.6, finetuned_weights=False, test=False, use_tta=False, normalize=True, quantiles=[0.0, 99.5]):
+        if device is None:
+            device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+        self.device = device
+        self.test = test
+        self.model = model
+        self.normalize = normalize
+        self.quantiles = quantiles
+
+        # Paths to model weights
+        self.model1_path, self.model2_path = get_weights(finetuned_weights)
+
+        # Load main models
+        self.model1 = self.load_model(self.model1_path, device=self.device)
+        self.model2 = self.load_model(self.model2_path, device=self.device) if model == "ensemble" or model == "model2" else None
+        if self.test:
+            # Define input and output paths for running test
+            self.input_path = os.path.join(os.path.dirname(__file__), 'resources/images')
+            self.output_path = os.path.join(os.path.dirname(__file__), 'resources/MEDIAR/results')
+        else:
+            self.input_path = input_path
+            self.output_path = output_path
+
+        # If using a single model
+        if self.model == "model1":
+            self.predictor = Predictor(
+                model=self.model1,
+                device=self.device,
+                input_path=self.input_path,
+                output_path=self.output_path,
+                algo_params={"use_tta": use_tta}
+            )
+
+        # If using a single model
+        if self.model == "model2":
+            self.predictor = Predictor(
+                model=self.model2,
+                device=self.device,
+                input_path=self.input_path,
+                output_path=self.output_path,
+                algo_params={"use_tta": use_tta}
+            )
+
+        # If using two models
+        elif self.model == "ensemble":
+            self.predictor = EnsemblePredictor(
+                model=self.model1,  # Pass model1 as model
+                model_aux=self.model2,  # Pass model2 as model_aux
+                device=self.device,
+                input_path=self.input_path,
+                output_path=self.output_path,
+                algo_params={"use_tta": use_tta}
+            )
+
+        if self.test:
+            self.run_test()
+
+        if not self.model in ["model1", "model2", "ensemble"]:
+            raise ValueError("Invalid model type. Choose from 'model1', 'model2', or 'ensemble'.")
+
+    def load_model(self, model_path, device):
+        model_args = {
+            "classes": 3,
+            "decoder_channels": [1024, 512, 256, 128, 64],
+            "decoder_pab_channels": 256,
+            "encoder_name": 'mit_b5',
+            "in_channels": 3
+        }
+        model = MEDIARFormer(**model_args)
+        weights = torch.load(model_path, map_location=device)
+        model.load_state_dict(weights, strict=False)
+        model.to(device)
+        model.eval()
+        return model
+    
+    def display_image_and_mask(self, img, mask):
+
+        from .plot import generate_mask_random_cmap
+        """
+        Displays the normalized input image and the predicted mask side by side.
+        """
+        # If img is a tensor, convert it to NumPy for display
+        if isinstance(img, torch.Tensor):
+            img = img.cpu().numpy()
+
+        # If mask is a tensor, convert it to NumPy for display
+        if isinstance(mask, torch.Tensor):
+            mask = mask.cpu().numpy()
+
+        # Transpose the image to have (H, W, C) format for display if needed
+        if len(img.shape) == 3 and img.shape[0] == 3:
+            img = img.transpose(1, 2, 0)
+
+        # Scale the normalized image back to [0, 255] for proper display
+        img_display = (img * 255).astype(np.uint8)
+
+        plt.figure(figsize=(10, 5))
+
+        # Display normalized image
+        plt.subplot(1, 2, 1)
+        plt.imshow(img_display)
+        plt.title("Normalized Image")
+        plt.axis("off")
+
+        r_cmap = generate_mask_random_cmap(mask)
+
+        # Display predicted mask
+        plt.subplot(1, 2, 2)
+        plt.imshow(mask, cmap=r_cmap)
+        plt.title("Predicted Mask")
+        plt.axis("off")
+
+        plt.tight_layout()
+        plt.show()
+
+    def predict_batch(self, imgs):
+        """
+        Predict masks for a batch of images.
+        
+        :param imgs: List of input images as NumPy arrays (each in (H, W, C) format).
+        :return: List of predicted masks as NumPy arrays.
+        """
+        processed_imgs = []
+
+        # Preprocess and normalize each image
+        for img in imgs:
+            if self.normalize:
+                # Normalize the image using the specified quantiles
+                img_normalized = normalize_image(img, lower_percentile=self.quantiles[0], upper_percentile=self.quantiles[1])
+            else:
+                img_normalized = img
+
+            # Convert image to tensor and send to device
+            img_tensor = torch.tensor(img_normalized.astype(np.float32).transpose(2, 0, 1)).to(self.device)  # (C, H, W)
+            processed_imgs.append(img_tensor)
+
+        # Stack all processed images into a batch tensor
+        batch_tensor = torch.stack(processed_imgs)
+
+        # Run inference to get predicted masks
+        pred_masks = self.predictor._inference(batch_tensor)
+
+        # Ensure pred_masks is always treated as a batch
+        if len(pred_masks.shape) == 3:  # If single image, add batch dimension
+            pred_masks = pred_masks.unsqueeze(0)
+
+        # Convert predictions to NumPy arrays and post-process each mask
+        predicted_masks = []
+        for pred_mask in pred_masks:
+            pred_mask_np = pred_mask.cpu().numpy()
+
+            # Extract dP and cellprob from pred_mask
+            dP = pred_mask_np[:2]  # First two channels as dP (displacement field)
+            cellprob = pred_mask_np[2]  # Third channel as cell probability
+
+            # Concatenate dP and cellprob along axis 0 to pass a single array
+            combined_pred_mask = np.concatenate([dP, np.expand_dims(cellprob, axis=0)], axis=0)
+
+            # Post-process the predicted mask
+            mask = self.predictor._post_process(combined_pred_mask)
+
+            # Append the processed mask to the list
+            predicted_masks.append(mask.astype(np.uint16))
+
+        return predicted_masks
+
+    def run_test(self):
+        """
+        Run the model on test images if the test flag is True.
+        """
+        # List of input images
+        imgs = []
+        img_names = []
+        
+        for img_file in os.listdir(self.input_path):
+            img_path = os.path.join(self.input_path, img_file)
+            img = io.imread(img_path)
+            
+            # Check if the image is grayscale (2D) or RGB (3D), and convert grayscale to RGB
+            if len(img.shape) == 2:  # Grayscale image (H, W)
+                img = np.repeat(img[:, :, np.newaxis], 3, axis=2)  # Convert grayscale to RGB
+
+            # Normalize the image if the normalize flag is True
+            if self.normalize:
+                img_normalized = normalize_image(img, lower_percentile=self.quantiles[0], upper_percentile=self.quantiles[1])
+            else:
+                img_normalized = img
+            
+            # Convert image to tensor and send directly to device
+            img_tensor = torch.tensor(img_normalized.astype(np.float32).transpose(2, 0, 1)).to(self.device)  # (C, H, W)
+
+            imgs.append(img_tensor)
+            img_names.append(os.path.splitext(img_file)[0])
+
+        # Stack all images into a batch (ensure it's always treated as a batch)
+        batch_tensor = torch.stack(imgs)
+
+        # Predict using the predictor (or ensemble predictor)
+        pred_masks = self.predictor._inference(batch_tensor)
+
+        # Ensure pred_masks is always treated as a batch
+        if len(pred_masks.shape) == 3:  # If single image, add batch dimension
+            pred_masks = pred_masks.unsqueeze(0)
+
+        # Convert predictions to NumPy arrays and post-process each mask
+        for i, pred_mask in enumerate(pred_masks):
+            # Ensure the dimensions of pred_mask remain consistent
+            pred_mask_np = pred_mask.cpu().numpy()
+
+            # Extract dP and cellprob from pred_mask
+            dP = pred_mask_np[:2]  # First two channels as dP (displacement field)
+            cellprob = pred_mask_np[2]  # Third channel as cell probability
+
+            # Concatenate dP and cellprob along axis 0 to pass a single array
+            combined_pred_mask = np.concatenate([dP, np.expand_dims(cellprob, axis=0)], axis=0)
+
+            # Post-process the predicted mask
+            mask = self.predictor._post_process(combined_pred_mask)
+
+            # Convert the mask to 16-bit format (ensure values fit into 16-bit range)
+            mask_to_save = mask.astype(np.uint16)
+
+            # Save the post-processed mask as a .tif file using cv2
+            mask_output_path = os.path.join(self.output_path, f"{img_names[i]}_mask.tiff")
+            cv2.imwrite(mask_output_path, mask_to_save)
+
+            print(f"Predicted mask saved at: {mask_output_path}")
+
+            self.display_image_and_mask(imgs[i].cpu().numpy(), mask)
+
+        print(f"Test predictions saved in {self.output_path}")
+
+    def preprocess_image(self, img):
+        """
+        Preprocess input image (numpy array) for compatibility with the model.
+        """
+        if isinstance(img, np.ndarray):  # Check if the input is a numpy array
+            if len(img.shape) == 2:  # Grayscale image (H, W)
+                img = np.repeat(img[:, :, np.newaxis], 3, axis=2)
+
+            elif img.shape[2] == 1:  # Single channel grayscale (H, W, 1)
+                img = np.repeat(img, 3, axis=2)  # Convert to 3-channel RGB
+
+            img_tensor = torch.tensor(img.astype(np.float32).transpose(2, 0, 1))  # Change shape to (C, H, W)
+        else:
+            img_tensor = img  # If it's already a tensor, assume it's in (C, H, W) format
+
+        return img_tensor.float()

spacr/plot.py

@@ -123,7 +123,7 @@ def plot_image_mask_overlay(file, channels, cell_channel, nucleus_channel, patho
 
     fig = _plot_merged_plot(image=image, outlines=outlines, outline_colors=outline_colors, figuresize=figuresize, thickness=thickness)
 
-    return
+    return fig
 
 def plot_masks(batch, masks, flows, cmap='inferno', figuresize=10, nr=1, file_type='.npz', print_object_number=True):
     """
@@ -274,6 +274,7 @@ def _generate_mask_random_cmap(mask):
     return random_cmap
 
 def _get_colours_merged(outline_color):
+    
     """
     Get the merged outline colors based on the specified outline color format.
 
@@ -283,6 +284,7 @@ def _get_colours_merged(outline_color):
     Returns:
     list: A list of merged outline colors based on the specified format.
     """
+
     if outline_color == 'rgb':
         outline_colors = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]  # rgb
     elif outline_color == 'bgr':
@@ -296,6 +298,7 @@ def _get_colours_merged(outline_color):
     return outline_colors
 
 def plot_images_and_arrays(folders, lower_percentile=1, upper_percentile=99, threshold=1000, extensions=['.npy', '.tif', '.tiff', '.png'], overlay=False, max_nr=None, randomize=True):
+    
     """
     Plot images and arrays from the given folders.
 
@@ -1055,7 +1058,7 @@ def _plot_recruitment(df, df_type, channel_of_interest, target, columns=[], figu
     plt.show()
 
     columns = columns + ['pathogen_cytoplasm_mean_mean', 'pathogen_cytoplasm_q75_mean', 'pathogen_periphery_cytoplasm_mean_mean', 'pathogen_outside_cytoplasm_mean_mean', 'pathogen_outside_cytoplasm_q75_mean']
-    columns = columns + [f'pathogen_slope_channel_{channel_of_interest}', f'pathogen_cell_distance_channel_{channel_of_interest}', f'nucleus_cell_distance_channel_{channel_of_interest}']
+    #columns = columns + [f'pathogen_slope_channel_{channel_of_interest}', f'pathogen_cell_distance_channel_{channel_of_interest}', f'nucleus_cell_distance_channel_{channel_of_interest}']
 
     width = figuresize*2
     columns_per_row = math.ceil(len(columns) / 2)

spacr/resources/MEDIAR/.git

@@ -0,0 +1 @@
+gitdir: ../../../.git/modules/spacr/resources/MEDIAR

spacr/resources/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

spacr/resources/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.

spacr/resources/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 CPU@2.26GHz                          |
+| 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}
+}
+```
+

spacr/resources/MEDIAR/SetupDict.py

@@ -0,0 +1,39 @@
+import torch.optim as optim
+import torch.optim.lr_scheduler as lr_scheduler
+import monai
+
+import core
+from train_tools import models
+from 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,
+}