Speed up segmentation and feature calculation

features.py

@@ -7,6 +7,7 @@
 import scipy.ndimage
 from PIL import Image
 from tqdm import tqdm
+from scipy.sparse import coo_matrix, csr_matrix
 
 Image.MAX_IMAGE_PIXELS = None
 
@@ -145,3 +146,106 @@ def df(self):
             self.compute()
         df = pd.DataFrame(self.computed_features, columns=self.feature_names)
         return df
+
+class FastNucleiFeatures:
+    """
+    Class for fast nuclear feuture calculations using sparse matrices. Memory intesive but speed up factor is ~1000x
+    """
+    def __init__(self, seg, orig, features='all', x_min=0, y_min=0, fast=True):
+        """
+        Class constructor, takes seg (predicted enumerated nuclei tiff) and orig (original image for colour calcluations)
+        """
+
+        self.seg = seg
+        self.orig = orig
+        self.coo = coo_matrix(seg)
+        self.computed_features = None
+        self.fast = fast
+        if self.fast:
+            self.indices = {i+1:ind for i,ind in enumerate(self.get_indices_sparse(self.seg))}
+        self.feature_dict = {'position': (self.position, ['x', 'y']),
+                             'size': (self.size, ['size']),
+                             'ellipse': (
+                                 self.ellips, ['first_axis', 'second_axis', 'ellipse_x', 'ellipse_y', 'ellipse_angle']),
+                             'color': (
+                                 self.color,
+                                 ['Blue_mean', 'Red_mean', 'Green_mean', 'Blue_std', 'Red_std', 'Green_std'])
+                             }
+        if features == 'all':
+            self.features = self.feature_dict.keys()
+        else:
+            self.features = features
+        self.x_min = x_min
+        self.y_min = y_min
+
+    def position(self, img, orig, **kwargs):
+        x, y = scipy.ndimage.measurements.center_of_mass(img)
+        x = x + self.x_min + kwargs['x_nucl']
+        y = y + self.y_min + kwargs['y_nucl']
+        return [x, y]
+
+    def ellips(self, img, orig, **kwargs):
+        try:
+            cont = cv.findContours(img.astype(np.uint8).T.copy(), cv.RETR_EXTERNAL, cv.CHAIN_APPROX_NONE)[1][0][:, 0, :]
+            ellipse_center, axles, angle = cv.fitEllipse(cont)
+            x, y = ellipse_center
+            x = x + self.x_min + kwargs['x_nucl'] 
+            y = y + self.y_min + kwargs['y_nucl']
+            if axles[1] > 100:
+                axles = (30, 30)
+            return [*axles, x, y, angle]
+        except:
+            return [0] * 5
+
+    def size(self, img, orig, **kwargs):
+        return [img.sum()]
+
+    def color(self, img, orig, **kwargs):
+        cell_pixels = orig[img]
+        return [*cell_pixels.mean(axis=0), *cell_pixels.std(axis=0)]
+
+    @property
+    def feature_names(self):
+        names = []
+        for f in self.features:
+            names += self.feature_dict[f][1]
+        return names
+
+    def compute_M(self, data):
+        cols = np.arange(data.size)
+        raveled = data.ravel()
+        cols = cols[raveled != 0]
+        raveled = raveled[raveled != 0]
+        return csr_matrix((cols, (raveled, cols)),
+                          shape=(data.max()+1, data.size))
+
+    def get_indices_sparse(self, data):
+        M = self.compute_M(data)
+        return [np.unravel_index(row.data, data.shape) for row in M[1:]]
+
+    def compute(self):
+        self.computed_features = []
+        img = self.seg
+        for i in tqdm.tqdm(range(1, img.max())):
+            if self.fast:
+                sub = self.indices[i] 
+            else:
+                coo_id = self.coo.data == i
+                sub = np.stack((self.coo.row[coo_id], self.coo.col[coo_id]))
+
+            upright = np.max(sub, axis=1)
+            downleft = np.min(sub, axis=1)
+            cut = img[downleft[0]:(upright[0]+1), downleft[1]:(upright[1]+1)]
+            orig_cut = self.orig[downleft[0]:(upright[0]+1), downleft[1]:(upright[1]+1)]
+            tmp_img = cut == i
+            tmp = []
+            for f in self.features:
+                tmp += self.feature_dict[f][0](tmp_img, orig_cut, x_nucl=downleft[0], y_nucl=downleft[1])
+            self.computed_features.append(tmp)
+        return self
+
+    def df(self):
+        if self.computed_features is None:
+            self.compute()
+        df = pd.DataFrame(self.computed_features, columns=self.feature_names)
+        return df