1#!python
2#cython: boundscheck=False
3#cython: wraparound=False
4#cython: cdivision=True
5import numpy as np
6from dipy.segment.mask import applymask
7from dipy.core.ndindex import ndindex
8from dipy.sims.voxel import add_noise
9cimport cython
10cimport numpy as cnp
11cdef extern from "dpy_math.h" nogil:
12    cdef double NPY_PI
13    cdef double NPY_INFINITY
14    double sqrt(double)
15    double log(double)
16    double exp(double)
17    double fabs(double)
18
19
20class ConstantObservationModel(object):
21    r"""
22    Observation model assuming that the intensity of each class is constant.
23    The model parameters are the means $\mu_{k}$ and variances $\sigma_{k}$
24    associated with each tissue class. According to this model, the observed
25    intensity at voxel $x$ is given by $I(x) = \mu_{k} + \eta_{k}$ where $k$
26    is the tissue class of voxel $x$, and $\eta_{k}$ is a Gaussian random
27    variable with zero mean and variance $\sigma_{k}^{2}$. The observation
28    model is responsible for computing the negative log-likelihood of
29    observing any given intensity $z$ at each voxel $x$ assuming the voxel
30    belongs to each class $k$. It also provides a default parameter
31    initialization.
32    """
33    def __init__(self):
34        r""" Initializes an instance of the ConstantObservationModel class
35        """
36        pass
37
38
39    def initialize_param_uniform(self, image, nclasses):
40        r""" Initializes the means and variances uniformly
41
42        The means are initialized uniformly along the dynamic range of
43        `image`. The variances are set to 1 for all classes
44
45        Parameters
46        ----------
47        image : array
48            3D structural image
49        nclasses : int
50            number of desired classes
51
52        Returns
53        -------
54        mu : array
55            1 x nclasses, mean for each class
56        sigma : array
57            1 x nclasses, standard deviation for each class.
58            Set up to 1.0 for all classes.
59        """
60        cdef:
61            double[:] mu = np.empty((nclasses,), dtype=np.float64)
62            double[:] sigma = np.empty((nclasses,), dtype=np.float64)
63
64        _initialize_param_uniform(image, mu, sigma)
65
66        return np.array(mu), np.array(sigma)
67
68
69    def seg_stats(self, input_image, seg_image, cnp.npy_intp nclass):
70        r""" Mean and standard variation for N desired  tissue classes
71
72        Parameters
73        ----------
74        input_image : ndarray
75            3D structural image
76        seg_image : ndarray
77            3D segmented image
78        nclass : int
79            number of classes (3 in most cases)
80
81        Returns
82        -------
83        mu, var : ndarrays
84            1 x nclasses dimension
85            Mean and variance for each class
86
87        """
88        cdef:
89            cnp.npy_intp i, j
90            cnp.npy_int16 s
91            double v
92            cnp.npy_intp size = input_image.size
93            double [::1] input_view
94            short [::1] seg_view
95            double [::1] mu_view
96            double [::1] var_view
97            cnp.npy_intp [::1] num_vox_view
98
99        # ravel C-contiguous array so we can use 1D indexing in the loop below
100        input_view = np.ascontiguousarray(input_image, dtype=np.double).ravel()
101        seg_view = np.ascontiguousarray(seg_image, dtype=np.int16).ravel()
102
103        mu = np.zeros(nclass, dtype=np.float64)
104        var = np.zeros(nclass, dtype=np.float64)
105        num_vox = np.zeros(nclass, dtype=np.intp)
106
107        mu_view = mu
108        var_view = var
109        num_vox_view = num_vox
110
111        for j in range(size):
112
113            s = seg_view[j]
114            v = input_view[j]
115
116            for i in range(nclass):
117
118                if s == i:
119                    mu_view[i] += v
120                    var_view[i] += v * v
121                    num_vox_view[i] += 1
122
123        mu = mu / num_vox
124        var = var / num_vox - mu * mu
125
126        return mu, var
127
128
129    def negloglikelihood(self, image, mu, sigmasq, cnp.npy_intp nclasses):
130        r""" Computes the gaussian negative log-likelihood of each class at
131        each voxel of `image` assuming a gaussian distribution with means and
132        variances given by `mu` and `sigmasq`, respectively (constant models
133        along the full volume). The negative log-likelihood will be written
134        in `nloglike`.
135
136        Parameters
137        ----------
138        image : ndarray
139            3D gray scale structural image
140        mu : ndarray
141            mean of each class
142        sigmasq : ndarray
143            variance of each class
144        nclasses : int
145            number of classes
146
147        Returns
148        -------
149        nloglike : ndarray
150            4D negloglikelihood for each class in each volume
151        """
152        cdef int l
153        nloglike = np.zeros(image.shape + (nclasses,), dtype=np.float64)
154
155        for l in range(nclasses):
156            _negloglikelihood(image, mu, sigmasq, l, nloglike)
157
158        return nloglike
159
160
161    def prob_image(self, img, cnp.npy_intp nclasses, mu, sigmasq, P_L_N):
162        r""" Conditional probability of the label given the image
163
164        Parameters
165        -----------
166        img : ndarray
167            3D structural gray-scale image
168        nclasses : int
169            number of tissue classes
170        mu : ndarray
171            1 x nclasses, current estimate of the mean of each tissue class
172        sigmasq : ndarray
173            1 x nclasses, current estimate of the variance of each
174            tissue class
175        P_L_N : ndarray
176            4D probability map of the label given the neighborhood.
177
178        Previously computed by function prob_neighborhood
179
180        Returns
181        --------
182        P_L_Y : ndarray
183            4D probability of the label given the input image
184        """
185        cdef int l
186        P_L_Y = np.zeros_like(P_L_N)
187        P_L_Y_norm = np.zeros_like(img)
188
189        g = np.empty_like(img)
190
191        for l in range(nclasses):
192            _prob_image(img, g, mu, sigmasq, l, P_L_N, P_L_Y)
193            P_L_Y_norm[:, :, :] += P_L_Y[:, :, :, l]
194
195        for l in range(nclasses):
196            P_L_Y[:, :, :, l] = P_L_Y[:, :, :, l] / P_L_Y_norm
197
198        return P_L_Y
199
200
201    def update_param(self, image, P_L_Y, mu, cnp.npy_intp nclasses):
202        r""" Updates the means and the variances in each iteration for all
203        the labels. This is for equations 25 and 26 of Zhang et. al.,
204        IEEE Trans. Med. Imag, Vol. 20, No. 1, Jan 2001.
205
206        Parameters
207        -----------
208        image : ndarray
209            3D structural gray-scale image
210        P_L_Y : ndarray
211            4D probability map of the label given the input image
212            computed by the expectation maximization (EM) algorithm
213        mu : ndarray
214            1 x nclasses, current estimate of the mean of each tissue
215            class.
216        nclasses : int
217            number of tissue classes
218
219        Returns
220        --------
221        mu_upd : ndarray
222                1 x nclasses, updated mean of each tissue class
223        var_upd : ndarray
224                1 x nclasses, updated variance of each tissue class
225        """
226        cdef int l
227        mu_upd = np.zeros(nclasses, dtype=np.float64)
228        var_upd = np.zeros(nclasses, dtype=np.float64)
229        mu_num = np.zeros(image.shape + (nclasses,), dtype=np.float64)
230        var_num = np.zeros(image.shape + (nclasses,), dtype=np.float64)
231
232        for l in range(nclasses):
233            mu_num[..., l] = P_L_Y[..., l] * image
234            var_num[..., l] = P_L_Y[..., l] * ((image - mu[l]) ** 2)
235
236            mu_upd[l] = np.sum(mu_num[..., l]) / np.sum(P_L_Y[..., l])
237            var_upd[l] = np.sum(var_num[..., l]) / np.sum(P_L_Y[..., l])
238
239        return mu_upd, var_upd
240
241
242cdef void _initialize_param_uniform(double[:, :, :] image, double[:] mu,
243                                    double[:] var) nogil:
244    r""" Initializes the means and standard deviations uniformly
245
246    The means are initialized uniformly along the dynamic range of `image`.
247    The standard deviations are set to 1 for all classes.
248
249    Parameters
250    ----------
251    image : array
252        3D structural gray-scale image
253    mu : array
254        buffer array for the mean of each tissue class
255    var : array
256        buffer array for the variance of each tissue class
257
258    Returns
259    -------
260    mu : array
261        1 x nclasses, mean of each class
262    var : array
263        1 x nclasses, variance of each class
264    """
265    cdef:
266        cnp.npy_intp nx = image.shape[0]
267        cnp.npy_intp ny = image.shape[1]
268        cnp.npy_intp nz = image.shape[2]
269        cnp.npy_intp nclasses = mu.shape[0]
270        int i
271        double min_val
272        double max_val
273    min_val = image[0, 0, 0]
274    max_val = image[0, 0, 0]
275    for x in range(nx):
276        for y in range(ny):
277            for z in range(nz):
278                if image[x, y, z] < min_val:
279                    min_val = image[x, y, z]
280                if image[x, y, z] > max_val:
281                    max_val = image[x, y, z]
282    for i in range(nclasses):
283        var[i] = 1.0
284        mu[i] = min_val + i * (max_val - min_val)/nclasses
285
286
287cdef void _negloglikelihood(double[:, :, :] image, double[:] mu,
288                            double[:] sigmasq, int classid,
289                            double[:, :, :, :] neglogl) nogil:
290    r""" Computes the gaussian negative log-likelihood of each class at
291    each voxel of `image` assuming a gaussian distribution with means and
292    variances given by `mu` and `sigmasq`, respectively (constant models
293    along the full volume). The negative log-likelihood will be written
294    in `neglogl`.
295
296    Parameters
297    ----------
298    image : array
299        3D structural gray-scale image
300    mu : array
301        mean of each class
302    sigmasq : array
303        variance of each class
304    classid : int
305        class identifier
306    neglogl : buffer for the neg-loglikelihood
307
308    Returns
309    -------
310    neglogl : array
311        neg-loglikelihood for the class (l = classid)
312    """
313    cdef:
314        cnp.npy_intp nx = image.shape[0]
315        cnp.npy_intp ny = image.shape[1]
316        cnp.npy_intp nz = image.shape[2]
317        cnp.npy_intp l = classid
318        cnp.npy_intp x, y, z
319        double eps = 1e-8      # We assume images normalized to 0-1
320        double eps_sq = 1e-16  # Maximum precision for double.
321
322    for x in range(nx):
323        for y in range(ny):
324            for z in range(nz):
325
326                if sigmasq[l] < eps_sq:
327
328                    if fabs(image[x, y, z] - mu[l]) < eps:
329                        neglogl[x, y, z, l] = 1 + log(sqrt(2.0 * NPY_PI *
330                                                           sigmasq[l]))
331                    else:
332                        neglogl[x, y, z, l] = NPY_INFINITY
333
334                else:
335                    neglogl[x, y, z, l] = (((image[x, y, z] - mu[l])**2.0) /
336                                           (2.0 * sigmasq[l]))
337                    neglogl[x, y, z, l] += log(sqrt(2.0 * NPY_PI *
338                                                    sigmasq[l]))
339
340
341cdef void _prob_image(double[:, :, :] image, double[:, :, :] gaussian,
342                      double[:] mu, double[:] sigmasq, int classid,
343                      double[:, :, :, :] P_L_N,
344                      double[:, :, :, :] P_L_Y) nogil:
345    r""" Conditional probability of the label given the image
346
347    Parameters
348    -----------
349    image : array
350        3D structural gray-scale image
351    gaussian : array
352        3D buffer for the gaussian distribution that is multiplied by
353        P_L_N to make P_L_Y
354    mu : array
355        current estimate of the mean of each tissue class
356    sigmasq : array
357        current estimate of the variance of each tissue
358        class
359    classid : int
360        tissue class identifier
361    P_L_N : array
362        4D probability map of the label given the neighborhood.
363        Previously computed by function prob_neighborhood
364    P_L_Y : array
365        4D buffer to hold P(L|Y)
366
367    Returns
368    --------
369    P_L_Y : array
370        4D probability of the label given the input image P(L|Y)
371    """
372    cdef:
373        cnp.npy_intp nx = image.shape[0]
374        cnp.npy_intp ny = image.shape[1]
375        cnp.npy_intp nz = image.shape[2]
376        cnp.npy_intp l = classid
377        cnp.npy_intp x, y, z
378
379        double eps = 1e-8
380        double eps_sq = 1e-16
381
382    for x in range(nx):
383        for y in range(ny):
384            for z in range(nz):
385
386                if sigmasq[l] < eps_sq:
387                    if fabs(image[x, y, z] - mu[l]) < eps:
388                        gaussian[x, y, z] = 1
389                    else:
390                        gaussian[x, y, z] = 0
391                else:
392                    gaussian[x, y, z] = (
393                        (exp(-((image[x, y, z] - mu[l]) ** 2) /
394                        (2 * sigmasq[l]))) / (sqrt(2 * NPY_PI * sigmasq[l])))
395
396                P_L_Y[x, y, z, l] = gaussian[x, y, z] * P_L_N[x, y, z, l]
397
398
399class IteratedConditionalModes(object):
400    def __init__(self):
401        pass
402
403    def initialize_maximum_likelihood(self, nloglike):
404        r""" Initializes the segmentation of an image with given
405            neg-loglikelihood
406
407        Initializes the segmentation of an image with neglog-likelihood field
408        given by `nloglike`. The class of each voxel is selected as the one
409        with the minimum neglog-likelihood (i.e. maximum-likelihood
410        segmentation).
411
412        Parameters
413        ----------
414        nloglike : ndarray
415            4D shape, nloglike[x, y, z, k] is the likelihhood of class k
416            for voxel (x, y, z)
417
418        Returns
419        --------
420        seg : ndarray
421            3D initial segmentation
422        """
423        seg = np.zeros(nloglike.shape[:3]).astype(np.int16)
424
425        _initialize_maximum_likelihood(nloglike, seg)
426
427        return seg
428
429
430    def icm_ising(self, nloglike, beta, seg):
431        r""" Executes one iteration of the ICM algorithm for MRF MAP
432        estimation. The prior distribution of the MRF is a Gibbs
433        distribution with the Potts/Ising model with parameter `beta`:
434
435        https://en.wikipedia.org/wiki/Potts_model
436
437        Parameters
438        ----------
439        nloglike : ndarray
440            4D shape, nloglike[x, y, z, k] is the negative log likelihood
441            of class k at voxel (x, y, z)
442        beta : float
443            positive scalar, it is the parameter of the Potts/Ising
444            model. Determines the smoothness of the output segmentation.
445        seg : ndarray
446            3D initial segmentation. This segmentation will change by one
447            iteration of the ICM algorithm
448
449        Returns
450        -------
451        new_seg : ndarray
452            3D final segmentation
453        energy : ndarray
454            3D final energy
455        """
456        energy = np.zeros(nloglike.shape[:3]).astype(np.float64)
457
458        new_seg = np.zeros_like(seg)
459
460        _icm_ising(nloglike, beta, seg, energy, new_seg)
461
462        return new_seg, energy
463
464
465    def prob_neighborhood(self, seg, beta, cnp.npy_intp nclasses):
466        r""" Conditional probability of the label given the neighborhood
467        Equation 2.18 of the Stan Z. Li book (Stan Z. Li, Markov Random Field
468        Modeling in Image Analysis, 3rd ed., Advances in Pattern Recognition
469        Series, Springer Verlag 2009.)
470
471        Parameters
472        -----------
473        seg : ndarray
474            3D tissue segmentation derived from the ICM model
475        beta : float
476            scalar that determines the importance of the neighborhood and
477            the spatial smoothness of the segmentation.
478            Usually between 0 to 0.5
479        nclasses : int
480            number of tissue classes
481
482        Returns
483        --------
484        PLN : ndarray
485            4D probability map of the label given the neighborhood of the
486            voxel.
487        """
488        cdef:
489            double[:, :, :] P_L_N = np.zeros(seg.shape, dtype=np.float64)
490            cnp.npy_intp classid = 0
491
492        PLN_norm = np.zeros(seg.shape, dtype=np.float64)
493        PLN = np.zeros(seg.shape + (nclasses,), dtype=np.float64)
494
495        for classid in range(nclasses):
496
497            P_L_N = np.zeros(seg.shape, dtype=np.float64)
498            _prob_class_given_neighb(seg, beta, classid, P_L_N)
499
500            PLN[:, :, :, classid] = np.array(P_L_N)
501            PLN[:, :, :, classid] = np.exp(- PLN[:, :, :, classid])
502            PLN_norm += PLN[:, :, :, classid]
503
504        for l in range(nclasses):
505            PLN[:, :, :, l] = PLN[:, :, :, l] / PLN_norm
506
507        return PLN
508
509
510cdef void _initialize_maximum_likelihood(double[:,:,:,:] nloglike,
511                                         cnp.npy_short[:,:,:] seg) nogil:
512    r""" Initializes the segmentation of an image with given
513    neg-log-likelihood.
514
515    Initializes the segmentation of an image with neg-log-likelihood field
516    given by `nloglike`. The class of each voxel is selected as the one with
517    the minimum neg-log-likelihood (i.e. the maximum-likelihood
518    segmentation).
519
520    Parameters
521    ----------
522    nloglike : array
523        4D nloglike[x, y, z, k] is the likelihhood of class k for voxel
524        (x, y, z)
525    seg : array
526        3D buffer for the initial segmentation
527
528    Returns :
529    seg : array,
530        3D initial segmentation
531    """
532    cdef:
533        cnp.npy_intp nx = nloglike.shape[0]
534        cnp.npy_intp ny = nloglike.shape[1]
535        cnp.npy_intp nz = nloglike.shape[2]
536        cnp.npy_intp nclasses = nloglike.shape[3]
537        double min_energy
538        cnp.npy_short best_class
539
540    for x in range(nx):
541        for y in range(ny):
542            for z in range(nz):
543
544                best_class = -1
545                for k in range(nclasses):
546                    if (best_class == -1) or (nloglike[x, y, z, k] <
547                                              min_energy):
548                        best_class = k
549                        min_energy = nloglike[x, y, z, k]
550                seg[x, y, z] = best_class
551
552
553cdef void _icm_ising(double[:,:,:,:] nloglike, double beta,
554                     cnp.npy_short[:,:,:] seg, double[:,:,:] energy,
555                     cnp.npy_short[:,:,:] new_seg) nogil:
556    r""" Executes one iteration of the ICM algorithm for MRF MAP estimation
557    The prior distribution of the MRF is a Gibbs distribution with the
558    Potts/Ising model with parameter `beta`:
559
560    https://en.wikipedia.org/wiki/Potts_model
561
562    Parameters
563    ----------
564    nloglike : array
565        4D nloglike[x, y, z, k] is the negative log likelihood of class k
566        at voxel (x, y, z)
567    beta : float
568        positive scalar, it is the parameter of the Potts/Ising model.
569        Determines the smoothness of the output segmentation
570    seg : array
571        3D initial segmentation.
572        This segmentation will change by one iteration of the ICM algorithm
573    energy : array
574        3D buffer for the energy
575    new_seg : array
576        3D buffer for the final segmentation
577
578    Returns
579    -------
580    energy : array
581        3D map of the energy for every voxel
582    new_seg : array
583        3D new final segmentation (there is a new one after each
584        iteration).
585    """
586    cdef:
587        cnp.npy_intp nneigh = 6
588        cnp.npy_intp* dX = [-1, 0, 0, 0,  0, 1]
589        cnp.npy_intp* dY = [0, -1, 0, 1,  0, 0]
590        cnp.npy_intp* dZ = [0,  0, 1, 0, -1, 0]
591        cnp.npy_intp nx = nloglike.shape[0]
592        cnp.npy_intp ny = nloglike.shape[1]
593        cnp.npy_intp nz = nloglike.shape[2]
594        cnp.npy_intp nclasses = nloglike.shape[3]
595        cnp.npy_intp x, y, z, xx, yy, zz, i, j, k
596        double min_energy = NPY_INFINITY
597        double this_energy = NPY_INFINITY
598        cnp.npy_short best_class
599
600    for x in range(nx):
601        for y in range(ny):
602            for z in range(nz):
603
604                best_class = -1
605                min_energy = NPY_INFINITY
606
607                for k in range(nclasses):
608                    this_energy = nloglike[x, y, z, k]
609
610                    for i in range(nneigh):
611                        xx = x + dX[i]
612                        if((xx < 0) or (xx >= nx)):
613                            continue
614                        yy = y + dY[i]
615                        if((yy < 0) or (yy >= ny)):
616                            continue
617                        zz = z + dZ[i]
618                        if((zz < 0) or (zz >= nz)):
619                            continue
620
621                        if seg[xx, yy, zz] == k:
622                            this_energy -= beta
623                        else:
624                            this_energy += beta
625
626                    if this_energy < min_energy:
627
628                        min_energy = this_energy
629                        best_class = k
630
631                new_seg[x, y, z] = best_class
632                energy[x, y, z] = min_energy
633
634
635cdef void _prob_class_given_neighb(cnp.npy_short[:, :, :] seg, double beta,
636                                   int classid, double[:, :, :] P_L_N) nogil:
637    r""" Conditional probability of the label given the neighborhood
638    Equation 2.18 of the Stan Z. Li book.
639
640    Parameters
641    -----------
642    image : array
643        3D structural gray-scale image
644    seg : array
645        3D tissue segmentation derived from the ICM model
646    beta : float
647        scalar that determines the importance of the neighborhood and the
648        spatial smoothness of the segmentation. Usually between 0 to 0.5
649    classid : int
650        tissue class identifier
651    P_L_N : array
652        buffer array for P(L|N)
653
654    Returns
655    --------
656    P_L_N : array
657        3D map of the probability of the label (l) given the neighborhood
658        of the voxel P(L|N)
659    """
660    cdef:
661        cnp.npy_intp nx = seg.shape[0]
662        cnp.npy_intp ny = seg.shape[1]
663        cnp.npy_intp nz = seg.shape[2]
664        cnp.npy_intp nneigh = 6
665        cnp.npy_intp l = classid
666        cnp.npy_intp x, y, z, xx, yy, zz
667        double vox_prob
668        cnp.npy_intp* dX = [-1, 0, 0, 0,  0, 1]
669        cnp.npy_intp* dY = [0, -1, 0, 1,  0, 0]
670        cnp.npy_intp* dZ = [0,  0, 1, 0, -1, 0]
671
672    for x in range(nx):
673        for y in range(ny):
674            for z in range(nz):
675
676                vox_prob = 0
677
678                for i in range(nneigh):
679                    xx = x + dX[i]
680                    if((xx < 0) or (xx >= nx)):
681                        continue
682                    yy = y + dY[i]
683                    if((yy < 0) or (yy >= ny)):
684                        continue
685                    zz = z + dZ[i]
686                    if((zz < 0) or (zz >= nz)):
687                        continue
688
689                    if seg[xx, yy, zz] == l:
690                        vox_prob -= beta
691                    else:
692                        vox_prob += beta
693
694                P_L_N[x, y, z] = vox_prob
695