saim1309/Cell_Segmentation

1import torch2import numpy as np3import os, sys4from monai.inferers import sliding_window_inference5 6sys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), "../../")))7 8from BasePredictor import BasePredictor9from utils import compute_masks10 11__all__ = ["Predictor"]12 13 14class Predictor(BasePredictor):15    def __init__(16        self,17        model,18        device,19        input_path,20        output_path,21        make_submission=False,22        exp_name=None,23        algo_params=None,24    ):25        super(Predictor, self).__init__(26            model,27            device,28            input_path,29            output_path,30            make_submission,31            exp_name,32            algo_params,33        )34        self.hflip_tta = HorizontalFlip()35        self.vflip_tta = VerticalFlip()36 37    @torch.no_grad()38    def _inference(self, img_data):39        """Conduct model prediction"""40 41        img_data = img_data.to(self.device)42        img_base = img_data43        outputs_base = self._window_inference(img_base)44        outputs_base = outputs_base.cpu().squeeze()45        img_base.cpu()46 47        if not self.use_tta:48            pred_mask = outputs_base49            return pred_mask50 51        else:52            # HorizontalFlip TTA53            img_hflip = self.hflip_tta.apply_aug_image(img_data, apply=True)54            outputs_hflip = self._window_inference(img_hflip)55            outputs_hflip = self.hflip_tta.apply_deaug_mask(outputs_hflip, apply=True)56            outputs_hflip = outputs_hflip.cpu().squeeze()57            img_hflip = img_hflip.cpu()58 59            # VertricalFlip TTA60            img_vflip = self.vflip_tta.apply_aug_image(img_data, apply=True)61            outputs_vflip = self._window_inference(img_vflip)62            outputs_vflip = self.vflip_tta.apply_deaug_mask(outputs_vflip, apply=True)63            outputs_vflip = outputs_vflip.cpu().squeeze()64            img_vflip = img_vflip.cpu()65 66            # Merge Results67            pred_mask = torch.zeros_like(outputs_base)68            pred_mask[0] = (outputs_base[0] + outputs_hflip[0] - outputs_vflip[0]) / 369            pred_mask[1] = (outputs_base[1] - outputs_hflip[1] + outputs_vflip[1]) / 370            pred_mask[2] = (outputs_base[2] + outputs_hflip[2] + outputs_vflip[2]) / 371 72        return pred_mask73 74    def _window_inference(self, img_data, aux=False):75        """Inference on RoI-sized window"""76        outputs = sliding_window_inference(77            img_data,78            roi_size=512,79            sw_batch_size=4,80            predictor=self.model if not aux else self.model_aux,81            padding_mode="constant",82            mode="gaussian",83            overlap=0.6,84        )85 86        return outputs87 88    def _post_process(self, pred_mask):89        """Generate cell instance masks."""90        dP, cellprob = pred_mask[:2], self._sigmoid(pred_mask[-1])91        H, W = pred_mask.shape[-2], pred_mask.shape[-1]92 93        if np.prod(H * W) < (5000 * 5000):94            pred_mask = compute_masks(95                dP,96                cellprob,97                use_gpu=True,98                flow_threshold=0.4,99                device=self.device,100                cellprob_threshold=0.5,101            )[0]102 103        else:104            print("\n[Whole Slide] Grid Prediction starting...")105            roi_size = 2000106 107            # Get patch grid by roi_size108            if H % roi_size != 0:109                n_H = H // roi_size + 1110                new_H = roi_size * n_H111            else:112                n_H = H // roi_size113                new_H = H114 115            if W % roi_size != 0:116                n_W = W // roi_size + 1117                new_W = roi_size * n_W118            else:119                n_W = W // roi_size120                new_W = W121 122            # Allocate values on the grid123            pred_pad = np.zeros((new_H, new_W), dtype=np.uint32)124            dP_pad = np.zeros((2, new_H, new_W), dtype=np.float32)125            cellprob_pad = np.zeros((new_H, new_W), dtype=np.float32)126 127            dP_pad[:, :H, :W], cellprob_pad[:H, :W] = dP, cellprob128 129            for i in range(n_H):130                for j in range(n_W):131                    print("Pred on Grid (%d, %d) processing..." % (i, j))132                    dP_roi = dP_pad[133                        :,134                        roi_size * i : roi_size * (i + 1),135                        roi_size * j : roi_size * (j + 1),136                    ]137                    cellprob_roi = cellprob_pad[138                        roi_size * i : roi_size * (i + 1),139                        roi_size * j : roi_size * (j + 1),140                    ]141 142                    pred_mask = compute_masks(143                        dP_roi,144                        cellprob_roi,145                        use_gpu=True,146                        flow_threshold=0.4,147                        device=self.device,148                        cellprob_threshold=0.5,149                    )[0]150 151                    pred_pad[152                        roi_size * i : roi_size * (i + 1),153                        roi_size * j : roi_size * (j + 1),154                    ] = pred_mask155 156            pred_mask = pred_pad[:H, :W]157 158        return pred_mask159 160    def _sigmoid(self, z):161        return 1 / (1 + np.exp(-z))162 163 164"""165Adapted from the following references:166[1] https://github.com/qubvel/ttach/blob/master/ttach/transforms.py167 168"""169 170 171def hflip(x):172    """flip batch of images horizontally"""173    return x.flip(3)174 175 176def vflip(x):177    """flip batch of images vertically"""178    return x.flip(2)179 180 181class DualTransform:182    identity_param = None183 184    def __init__(185        self, name: str, params,186    ):187        self.params = params188        self.pname = name189 190    def apply_aug_image(self, image, *args, **params):191        raise NotImplementedError192 193    def apply_deaug_mask(self, mask, *args, **params):194        raise NotImplementedError195 196 197class HorizontalFlip(DualTransform):198    """Flip images horizontally (left -> right)"""199 200    identity_param = False201 202    def __init__(self):203        super().__init__("apply", [False, True])204 205    def apply_aug_image(self, image, apply=False, **kwargs):206        if apply:207            image = hflip(image)208        return image209 210    def apply_deaug_mask(self, mask, apply=False, **kwargs):211        if apply:212            mask = hflip(mask)213        return mask214 215 216class VerticalFlip(DualTransform):217    """Flip images vertically (up -> down)"""218 219    identity_param = False220 221    def __init__(self):222        super().__init__("apply", [False, True])223 224    def apply_aug_image(self, image, apply=False, **kwargs):225        if apply:226            image = vflip(image)227 228        return image229 230    def apply_deaug_mask(self, mask, apply=False, **kwargs):231        if apply:232            mask = vflip(mask)233 234        return mask235