ericup/celldetection

1import torch2import celldetection as cd3import cv24import numpy as np5 6__all__ = ['contours2labels', 'CpnInterface']7 8 9def contours2labels(contours, size, overlap=False, max_iter=999):10    labels = cd.data.contours2labels(cd.asnumpy(contours), size, initial_depth=3)11 12    if not overlap:13        kernel = cv2.getStructuringElement(1, (3, 3))14        mask_sm = np.sum(labels > 0, axis=-1)15        mask = mask_sm > 1  # all overlaps16        if mask.any():17            mask_ = mask_sm == 1  # all cores18            lbl = np.zeros(labels.shape[:2], dtype='float64')19            lbl[mask_] = labels.max(-1)[mask_]20            for _ in range(max_iter):21                lbl_ = np.copy(lbl)22                m = mask & (lbl <= 0)23                if not np.any(m):24                    break25                lbl[m] = cv2.dilate(lbl, kernel=kernel)[m]26                if np.allclose(lbl_, lbl):27                    break28        else:29            lbl = labels.max(-1)30        labels = lbl.astype('int')31    return labels32 33 34class CpnInterface:35    def __init__(self, model, device=None, **kwargs):36        self.device = ('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device37        model = cd.resolve_model(model, **kwargs)38        if not isinstance(model, cd.models.LitCpn):39            model = cd.models.LitCpn(model)40        self.model = model.to(device)41        self.model.eval()42        self.model.requires_grad_(False)43        self.tile_size = 166444        self.overlap = 38445 46    def __call__(47            self,48            img,49            div=255,50            reduce_labels=True,51            return_labels=True,52            return_viewable_contours=True,53    ):54        if img.ndim == 2:55            img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)56        img = img / div57        x = cd.data.to_tensor(img, transpose=True, dtype=torch.float32)[None]58        with torch.no_grad():59            out = cd.asnumpy(self.model(x, crop_size=self.tile_size,60                                        stride=max(64, self.tile_size - self.overlap)))61            # if torch.cuda.device_count():62            #     print(cd.GpuStats())63 64        contours, = out['contours']65        boxes, = out['boxes']66        scores, = out['scores']67 68        labels = None69        if return_labels or return_viewable_contours:70            labels = contours2labels(contours, img.shape[:2], overlap=not reduce_labels)71 72        return dict(73            contours=contours,74            labels=labels,75            boxes=boxes,76            scores=scores77        )78