PROBABLISTIC SEGMENTATION

Abstract

A machine learning system may be used for determining if a segmentation of a medical image is a reasonable segmentation in the sense that it is a segmentation that could be made by a human user and does not contain any impossible combinations of pixel values. The method is enhanced by user input to avoid the impossible combinations.

Claims

1. A computer-implemented method of determining an improved segmentation of a medical image including a set X of pixels or voxels having a given pixel values x to a corresponding set of Y of segmentation values y, using a machine learning system approximating the unnormalized conditional probability distribution f(x, y)≈C.Math.p(y|x), with a neural network, the method comprising the steps of a. providing the medical image to a neural network based machine learning system that has been trained by means of training data sets comprising manually made segmentations, to determine if a segmentation is similar to the segmentations in the training dataset, obtaining an initial segmentation of the medical image from the machine learning system in the form of a set Y of segmentation values y, said segmentation values y indicating for each pixel if it belongs to the structure, b. manually adjusting the initial segmentation by changing at least one of the segmentation values y to y.sub.1 produce an adjusted set Y′ of segmentation values y′=[y.sub.1, y.sub.2], and c. updating, by the machine learning system, the initial segmentation to produce an updated segmentation by optimizing the set Y of segmentation values based on the given pixel values x and the adjusted set Y′ of segmentation values, by solving the optimization problem y.sub.2*=argmax.sub.y.sub.2.sub.∈[0,1].sub.Mf(x, [y.sub.1, y.sub.2]).

2. A method according to claim 1, further comprising repeating steps b and c using the updated segmentation as the initial segmentation, until the updated segmentation is found to be good enough.

3. A method according to claim 1, wherein the optimization problem y.sub.2*=argmax.sub.y.sub.2.sub.∈[0,1].sub.Mf(x, [y.sub.1, y.sub.2]) is solved by means of simulated annealing.

4. A method according to claim 1, wherein the optimization problem y.sub.2*=argmax.sub.y.sub.2.sub.∈[0,1].sub.Mf(x, [y.sub.1, y.sub.2]) is solved by means of continuous approximation of the basic problem.

5. A method according to claim 1, wherein the machine learning system is a generative adversarial network—GAN—and the optimization is performed using a discriminator that has been trained to recognize a segmentation that are similar to those in a training dataset of segmentations that are or resemble manually made segmentations.

6. A method according to claim 1, wherein the machine learning system is an auto encoder and the optimization is performed using an encoder and a and the optimization is based on the score f(x,y)=p.sub.Z(e(x,y)).

7. A medical planning method including segmenting a medical image according to claim 1 and using the segmented medical image as a basis for planning.

8. A computer program product for generating an improved segmentation of a medical image including a set X of pixels having given pixel values x using a neural network based machine learning system, the computer program product comprising computer readable code means which, when run in a computer will cause the computer to perform the following steps: obtaining an initial segmentation of a structure in the image in the form of a set Y of segmentation values y, said segmentation values y indicating for each pixel if it belongs to the structure, receiving manual input for adjusting the initial segmentation by changing at least one of the segmentation values y to y and produce an adjusted set Y′ of segmentation values y′=[y.sub.1, y.sub.2] based on the manual input updating, by the machine learning system the initial segmentation to produce an updated segmentation by optimizing the set Y of segmentation values based on the given pixel values X and the adjusted set Y′ of segmentation values by solving the optimization problem y.sub.2*=argmax.sub.y.sub.2.sub.∈[0,1].sub.Mf(x, [y.sub.1, y.sub.2]) optionally repeating steps b and c using the updated segmentation as the initial segmentation, until the updated segmentation is found to be good enough.

9. A computer system for generating a segmentation of a medical image, comprising a processor, a program memory arranged to hold at least one computer program to be run in the processor, and a data memory, wherein the program memory holds a computer program product according to claim 8.

10. A method of training a machine learning system capable of generating a score to a segmentation of at least one structure in a medical image, comprising: inputting training data sets into the machine learning system, each set comprising a medical image together with one or more manually made segmentations and/or one or more machine generated segmentations of the at least one structure and training the machine learning system to recognize segmentations that are similar to those in the training dataset.

11. A machine learning system trained according to the method of claim 10.

12. A machine learning system according to claim 11, which is an almost everywhere differential neural network.

Description

BRIEF DESCRIPTION OF THE DRAWINGS

[0031] The invention will be described in more detail in the following, by way of example and with reference to the appended drawings, in which

[0032] FIG. 1 is a flowchart of an overall method according to an embodiment of the invention;

[0033] FIGS. 2 and 3 are examples of possible segmentations with manual adjustments entered by a human;

[0034] FIG. 4 is a diagram showing the probability p for each of a number of possible segmentations of that segmentation being a plausible segmentation;

[0035] FIGS. 5 and 6 illustrate different parts of an embodiment of the method using a GAN;

[0036] FIG. 7 illustrates an embodiment of the method using an autoencoder;

[0037] FIG. 8 illustrates a computer system in which the inventive method may be implemented.

DETAILED DESCRIPTION

[0038] FIG. 1 is an overall flowchart of a method according to an embodiment of the invention. The method uses a machine learning system that is able to generate, based on an input set comprising a CT image and a possible segmentation of the CT image, a probability indicating how likely it is that a doctor would generate the segmentation. Before the system can be used to generate such a value, it must be trained based on a number of CT images, each with one or several segmentations associated with it. Various machine learning systems, and how to train them, are known to the skilled person. Some examples will be discussed in some detail below.

[0039] In step S51 a segmentation of a structure in a CT image is obtained, in any suitable way. The segmentation can be completely arbitrary. In step S52, a doctor considers the segmentation and makes a manual adjustment to a part of the structure. Some possible reasons for adjustment, and types of adjustment, will be discussed below. In step S53, the discriminator is used to update the segmentation to identify a segmentation that is deemed to be a plausible segmentation, in that there is a high likelihood that it could be arrived at by a doctor. This is performed as an optimization, which will be discussed below. Step S54 is a decision step for the user, to determine if the updated segmentation is good enough. If yes, the segmentation is output and may be used, for example, for treatment planning. If no, the procedure returns to step S52 or S53 for further adjustment. The user is typically a medical practitioner, depending on the procedure for which the segmentation is to be used.

[0040] In the following, the medical image that is to be segmented is denoted as x, the proposal segmentation for y and f(x,y) as our approximation of the unnormalized likelihood C.Math.p(y|x), where C is some positive unknown constant. The updating in step S53 may be performed as an optimization according to the following: Mathematically, f is a high dimensional functional with 2N parameters, where the first N parameters, given by x, are associated with the image and represented by real numbers and the remaining N parameters, given by N, are binary and associated with the segmentation. If no restrictions are made on the number of voxels considered in the image, N is typically of order 100.sup.3.

[0041] In the following, the portion of the segmentation y that has been entered by the user is denoted y.sub.1 and the portion of the segmentation that corresponds to the automatically generated segmentation is denoted y.sub.2. For notational convenience, without loss of generality, in the following, only the case when there is one object to delineate against the background will be considered. The optimization problem to be solved can be expressed as finding the y.sub.2 generating a maximal value of f for a given combination of x and y.sub.1, that is,

y.sub.2*=argmax.sub.y.sub.2.sub.∈[0,1].sub.Mf(x,[y.sub.1,y.sub.2]

As per convention, f(x,y) denotes a neural network that assigns a score to how well y matches x according to a specific group of clinical practitioners. Assuming that the user of the system has manually provided a part of y denoted y.sub.1 the rest of y may be denoted with y.sub.2. Based on this, a prediction of y.sub.2 may be generated by solving the argmax problem as defined above. This can be done in an iterative fashion, where a user starts by letting y.sub.1 be empty, and then add parts of to y.sub.1 and resolving the problem.

[0042] This optimization can be approached in many ways by methods known in the art. A simple solution might be simulated annealing, which is suitable for discrete problems without any simple structure. If all of the parameters for one regular segmentation is optimized over, this method will not be feasible. The size of the problem may however be reduced by excluding all voxels that are clearly either part of the segmented structure or not. In other words, the size may be reduced by excluding the voxels for which there is no or little uncertainty, resulting in a problem encompassing only the voxels for which it is uncertain if they belong to the structure or not. This can be defined by user input, or by a method identifying voxels which are uncertain. For example, the network can be a regular U-net that assigns an independent probability to each voxel and filter out the voxels that have less probability associated with every structure than a certain threshold probability.

[0043] In other words, the set X represents the pixels or voxels of a medical image and may be a real valued vector. The sets Y and Y′ are preferably vectors, such as binary segmentation vectors were ones indicate foreground and zeros indicate background.

[0044] An alternative method, which is able to handle larger problems, is by continuous approximation of the basic problem. That is, by letting yϵ[0,1].sup.N take values in the unit cube rather than just the space of binary vectors {0,1}.sup.N. This continuum can then be interpreted as either “the fraction of the voxel belonging to a structure” or “the probability of a voxel being in that structure”. Such an approximation is sometimes done naturally to capture that a voxel on the boundary may only be part of some structure to a certain fraction. A requirement for this approximation is that the function f can be evaluated over the unit cube of dimension N yϵ[0,1].sup.N rather than just the space of binary vectors of dimension N {0,1}N. Also, the function f needs to be almost everywhere differentiable with respect to y. If zϵR.sup.M is defined by y.sub.2=sigmoid(z), where the sigmoid function operates elementwise, a gradient based optimization method such as gradient descent may be used to solve the optimization problem

z*=argmax.sub.zf(x,[y.sub.1,sigmoid(z)])

without any constraints. A final segmentation proposal can then be obtained by taking y.sub.2*=sigmoid(z*) or y.sub.2*=round(sigmoid(z*)) depending on whether binary or continuous representations are used. Here round( ) rounds off each element to 0 or 1 based on which of them are closest.

[0045] Given that the machine learning system is a neural network, the backpropagation algorithm may be used as an efficient way to implement gradient descent in a neural network. Such neural networks can be achieved in different ways. Examples of how to use the discriminator associated with a GAN or the encoder associated with a specific kind of auto-encoder will be discussed below.

[0046] FIG. 2 shows an initial segmentation y of an organ, as a solid line 21. In this example, a user considers that an area y1 outside of the segmentation, shown delineated by a dashed line 22, also belongs to the structure. Bearing in mind that the image is in 3D and may include a large number of slices, delineating the whole structure again according to this will be a lot of work. As an example, a CT scan of a spine may comprise approximately 200 slices.

[0047] According to embodiments of the invention, the user can mark one or more pixels within the segmentation as “definitely part of the structure” or “not part of the structure” or “on the edge of the structure”. In this example, the user corrects the initial segmentation, marked by a solid line, by marking out some points along what they consider to be the edge of the structure. Alternatively, one or more points within the area y1 may be marked as belonging to the structure. The discriminator will then select among the possible segmentations that coincide with the user marking, one that has a high probability. This means that the discriminator will select a segmentation that looks reasonable, for example in that it includes a continuous structure incorporating adjacent pixels that have similar properties. The adjustments should be made in a suitable position in the structure to enable the discriminator to make the adjustment over the relevant part of the outline. For example, if only one adjusted point is entered, it should typically be in the middle of the portion of the outline that is to be corrected.

[0048] FIG. 3 shows another situation, illustrating the importance of considering the probability of the segmentation as a whole. In a) a possible initial segmentation y.sub.init of a structure is shown as a solid line. In this example, a user assessing the segmentation considers it is clear that everything within the smaller area y.sub.corr to the left of the dashed line belongs to the structure whereas the area y.sub.not between the dashed line and the solid line on the right is not part of the structure. The voxels in this area y.sub.not seem to have uniform properties. This whole area, therefore, should either be included or excluded in the segmentation of the structure. In this example, the doctor considers that it should be excluded. As for FIG. 2, delineating the part of the structure that needs to be changed in 3D manually involves a lot of work.

[0049] As in the previous example, the user can mark one or more pixels within the segmentation as “part of the structure” or “not part of the structure” or “on the edge of the structure”. b) shows three such points, p1, p2, p3 set by the doctor to identify the border of the structure. To facilitate the calculations the points are selected more or less evenly distributed over the portion of the outline that is to be adjusted. As can be seen in b), the resulting structure of adding these three points to the segmentation has a very unlikely shape for any structure, with three sharp dents into the initial segmentation y.sub.init, and would likely have the probability zero. According to this embodiment of the invention, therefore, our approximation of the unnormalized probability distribution C.Math.p(y|x), is used to identify the most likely segmentations that fulfil the requirements added by the user. The resulting segmentation will probably look like the one shown in c), that is, the whole area outside of the dashed line is considered not to be a part of the structure.

[0050] FIG. 4 is a simplified diagram in one dimension illustrating the unnormalized likelihood approximation f(x,y) as a function of y while keeping x fixed to one specific input image. Two segmentations, y.sub.1 and y.sub.2 are marked on the horizontal axis and as can be seen, the evaluation of the first segmentation f(x,y.sub.1) is much lower than f(x,y.sub.2). This means, according to the model, that it is much more likely that a doctor would manually have delineated the patient as y.sub.2 rather than y.sub.1.

[0051] Using the values represented in FIG. 4, one of the possible segmentations may be selected to be used as a starting point to arrive at a more plausible segmentation by means of a gradient method. The change in probability between the selected segmentation and the one immediately to the right and left of it is determined and the one that results in the greatest increase in probability is selected. The procedure then continues in the same direction until there is no longer an increase in the probability. In that way, a segmentation having a high probability should be reached, although as can be seen, not necessarily the segmentation having the highest possible probability. Starting from the first segmentation y.sub.1, therefore, the segmentation to the right of it should be selected repeatedly until there is no positive gradient on the curve, where the probability no longer increases. The segmentation selected by this procedure may be evaluated by a user, who may then make further adjustment to the selected segmentation to provide a new starting point. The procedure may then be repeated starting from the new starting point. In principle, any segmentation that coincides with the user specified corrections may be used as a starting point. For instance, one could fix the user specified variables and randomly draw the rest of the variables as independent U[0,1] variables. Another way would be to use the output of a classical U-net model to populate the free variables. That is, again, freeze the variables associated with user input and use the U-net output as starting points for the variables for which the optimization will be performed.

[0052] The machine learning system used according to the invention may be any type of system that is capable of generating a score reflecting the probability that a certain segmentation was generated by a doctor.

[0053] The machine learning system may be any suitable type of neural network, trained in any suitable way. For instance, one may use a part of a Generative Adversarial Network (GAN), known per se. Such networks were first described in Goodfellow, Ian, et al. “Generative adversarial nets.” Advances in neural information processing systems. 2014. With a GAN, two competing neural networks are used, a generator and a discriminator. During training, a number of training sets, each comprising a medical image and one or more segmentations of one or more structures in the medical image, are fed to the discriminator. Some segmentations are authentic segmentations of the medical image, typically created manually and others are generated by the generator. The generator strives to create pictures, or in the context of this invention, segmentations, that will be perceived as authentic, while the discriminator strives to differentiate between authentic and machine-made segmentations. In this way, the generator and the discriminator trigger each other to become better.

[0054] The function of the GAN generator can be expressed as

g.sub.w(medical image,noise)->segmentation of tumor and/or organ

The noise is randomly generated to provide a variation of segmentations when the generator is run several times.

[0055] The GAN generator works according to the equation

g.sub.w(x)=y

where g is a function subjected to one or more weights w, where the weights may be adjusted such that for each CT image x a binary image y will be generated in which each pixel is identified as either belonging or not belonging to the structure. The process of training a GAN generator is known to the skilled person.

[0056] In this way, the generator is trained to generate segmentations that are increasingly similar to authentic segmentations created by a human, while the discriminator is trained to make the distinction between authentic and machine-made segmentations with greater accuracy. Formally, the discriminator is trained to approximate the probability distribution

d(x,y)≈p(b=1|x,y)

where b is a variable taking the value 1 if a segmentation is done by a human and the value 0 otherwise. Since

p(b|x,y)∝p(y|x,b)

and b is fixed as b=1, d is proportional to the desired probability distribution, that is, the probability that a doctor would have delineated a medical image x with y. Since the discriminator is a neural network proportional to the likelihood p(y|x), the next step may be to find a partial segmentation as described above with f(x,y)=d(x, y).

[0057] FIG. 5 is a flow chart of a possible method, known per se, for automatic generation of segmentations using a GAN. In a first step S61, a CT image and some general noise is input to the generator. In step S62, the generator creates a segmentation based on the input data from step S61. The segmentation is output in step S63. This procedure may be repeated as many times as desired, if a set of segmentations is desired.

[0058] FIG. 6 is a flowchart of an embodiment of the inventive method performed by making use of the discriminator. In step S71 a CT image and a segmentation of the relevant part of the CT image is input to the discriminator. The segmentation may be obtained in any suitable way, for example manually or by the process of FIG. 1 or FIG. 5. In step S72, the discriminator generates a score S73 indicative of how likely it is, given the CT image, that the segmentation was made by a doctor. If necessary, the segmentation may be adjusted manually, in step S74, and submitted again to the discriminator. This is typically done if the likelihood is too low, that is, if the segmentation is too far from a segmentation that a doctor would generate, or if the doctor determines that manual adjustment is desirable. It is also possible to perform a manual adjustment like the one in step S74 before step S72.

[0059] The discriminator is arranged to generate an optimized segmentation by solving the problem argmax.sub.y f(x,y), f being the discriminator. The doctor will consider the output from the discriminator and, if necessary, perform an adjustment in step S74 and revert to step S72 to and generate a new optimized segmentation. That is, if a correction is provided, y will be split up into y_1 and y_2, where y_1 is the provided adjustments. Then argmax y.sub.2 f(x,[y.sub.1, y.sub.2]). This will then be repeated until the results are satisfactory to the user.

[0060] As an alternative, the machine learning model can be a version of an auto-encoder. Such networks, and how they may be set up, are known per se. One way to set them up may be described as follows, with reference to FIG. 7: The auto-encoder must be regularized in a way such that the data in the latent-space looks approximately p.sub.Z-distributed, schematically illustrated by the dashed curve in the rightmost coordinate system. Formally, consider again a dataset of n manually delineated medical images images (x.sub.i, y.sub.i)˜p.sub.X,Y, where an example of p.sub.X,Y is shown in the leftmost coordinate system. Let a neural network encoder be denoted with e∈E and a neural network decoder be denoted with d∈D, where E and D are two predefined sets of neural networks. The encoder is illustrated in FIG. 7 by an arrow denoted e(x,y) pointing from the leftmost to the rightmost coordinate system and the decoder is illustrated by an arrow denoted d(x,y) pointing in the opposite direction. Let L.sub.R be some reconstruction loss for both x and y and L.sub.Z be some loss associated with how close the empirical distribution of the datapoints in the latent space are to the target distribution p.sub.Z. This corresponds to a an average of dirac measures taking values in a couple samples from the distribution depicted with the solid curve p(e(x,y)) in the rightmost coordinate system. The network is then trained by solving

[00002]$\arg {\min }_{e\in E,d\in D}\frac{1}{n}\underset{i=1}{\overset{n}{.Math.}}{L}_R\left(d\left(e\left(x,y\right)\right),x,y\right)+L_Z\left(\frac{1}{n}\underset{i=1}{\overset{n}{.Math.}}\delta \left\{e\left(x_i,y_i\right)\right\}\right),$

where δ is the indicator function, and the resulting encoder is used to approximate the joint distribution f(x,y)=p.sub.Z(e(x, y))≈p.sub.X,Y(x, y), necessary for our method described above.

[0061] Furthermore, if f.sub.z is chosen to be almost everywhere differentiable, for instance normally distributed, this optimization problem described in above for updating segmentations given f, can be solved with a back-propagation scheme.

[0062] For example, let the network be a fully convolutional encoder and a fully convolutional decoder, let all the activation functions be ReLU-activations and let the output of the decoder that is associated with the segmentation, pass through a softmax activation as the last layer. The loss functions can be defined such that the reconstruction of x is given by a least squares loss and the reconstruction for y is given by the mean cross entropy loss. Finally, the regularizer loss L.sub.Z may be given by the Wasserstein loss. In order to train this model, stochastic gradient descent with mini batches can not be used in the standard way since the regularizer term is dependent on all of the datapoints in the dataset. To handle this, we approximate the regularizer term for a minibatch of the pair (x.sub.j, y.sub.j) with

[00003]$L_Z=\left(\frac{1}{n}\underset{i=1}{\overset{n}{.Math.}}\delta \left\{e\left(x_i,y_i\right)\right\}\right)\approx L_Z\left(\delta \left\{e\left(x_j,y_j\right)\right\}+\frac{1}{T-1}\underset{i=1}{\overset{T-1}{.Math.}}{\delta }_t\right)$

where δ.sub.0 is the first evaluated latent point and δ.sub.T-1 was the previously evaluated latent point during training. This specifies all of the free parameters that are included in the model.

[0063] FIG. 8 is a schematic representation of a computer system in which the inventive planning method may be performed. A computer 51 comprises a processor 53, a data memory 54 and a program memory 55. Preferably, a user input means 58 is also present, in the form of a keyboard, a mouse, a joystick, voice recognition means or any other available user input means.

[0064] The data memory 54 comprises data such as images and possible segmentation, and typically one or more objective functions. The data in the data memory may be generated in the computer 51, entered by means of the user input means or received from another storage means, in any way known in the art.

[0065] As will be understood, the data memory 54 is only shown schematically. There may be several data memory units, each holding one or more different types of data, for example, one data memory for the value set, one for the objective function, etc.

[0066] The program memory 55 holds a computer program arranged to control the processor to perform the inventive method. It will be understood that not all of the method steps are necessarily performed in the computer 51.