Method, device, and computer program product for self-supervised learning of pixel-wise anatomical embeddings in medical images

1. A self-supervised anatomical embedding (SAM) method for medical images, the method comprising:

randomly selecting a plurality of images from an unlabeled image batch;

for each image of the plurality of images,

performing random data augmentation to obtain a patch pair which is inputted to a neural network;

generating global and local embedding tensors for each patch of the patch pair using the neural network;

selecting positive pixel pairs from the patch pair and obtaining positive embedding pairs corresponding to the positive pixel pairs, wherein the positive embedding pairs include global positive embedding pairs and local positive embedding pairs;

for each positive pixel pair:

computing global and local similarity maps using the global and local embedding tensors;

finding global hard negative embeddings using the global similarity maps;

selecting global random negative embeddings from a plurality of patch pairs; and

pooling the global hard negative embeddings and the global random negative embeddings to obtain final global negative embeddings; and

finding local hard negative embeddings using the global and local similarity maps, and randomly sampling final local negative embeddings from the local hard negative embeddings; and

minimizing a final info noise contrastive estimation (InfoNCE) loss including global and local InfoNCE losses, wherein the global InfoNCE loss is computed using the global positive embedding pairs and the final global negative embeddings, and the local InfoNCE loss is computed using the local positive embedding pairs and the final local negative embeddings.

2. The method according to claim 1 , wherein the global and local InfoNCE losses are computed, respectively, according to:

L

=

-

∑

i

=

1

n

p

⁢

o

⁢

s

⁢

log

⁢

exp

⁡

(

f

i

·

f

i

′

/

τ

)

exp

⁡

(

f

i

·

f

i

′

/

τ

)

+

∑

j

=

1

n

n

⁢

e

⁢

g

⁢

exp

⁡

(

f

i

·

h

i

⁢

j

/

τ

)

wherein L denotes a InfoNCE loss, i denotes an i-th element of positive pixel pairs, n pos denotes a number of positive pixels, n neg denotes a number of negative pixels, f i denotes an i-th positive embedding in one of the positive embedding pairs, f i ′ denotes an i-th positive embedding in the other of the positive embedding pairs, τ denotes a temperature parameter, and “.” denotes an inner production operation.

3. The method according to claim 1 , wherein selecting the positive pixel pairs includes:

when one patch of the patch pair overlaps another patch of the patch pair, randomly selecting the positive pixel pairs from an overlapping area of the one patch of the patch pair and the another patch of the patch pair; and

when the one patch of the patch pair does not overlap the another patch of the patch pair, randomly sampling the positive pixel pairs from each patch of the patch pair.

4. The method according to claim 1 , wherein:

for each positive pixel pair, randomly selecting one or more negative pixels from the patch pair, wherein a distance between each of the one or more negative pixels and the each positive pixel pair is greater than a preset value.

5. The method according to claim 1 , wherein:

the SAM method generates semantic embeddings for each pixel which describes a corresponding anatomical location or body part.

6. The method according to claim 1 , wherein:

after a point of interest is labeled on a template image, a same body part in other images is located by simple nearest neighbor searching.

7. The method according to claim 2 , wherein the final InfoNCE loss is computed according to:

L final =L g +L l

wherein L final denotes the final InfoNCE loss, L g denotes the global InfoNCE loss, and L l denotes the local InfoNCE loss.

8. A device for self-supervised anatomical embedding (SAM), comprising:

a memory, containing a computer program stored thereon; and

a processor, coupled with the memory and configured, when the computer program being executed, to perform a method including:

randomly selecting a plurality of images from an unlabeled image batch;

for each image of the plurality of images,

performing random data augmentation to obtain a patch pair which is inputted to a neural network;

generating global and local embedding tensors for each patch of the patch pair using the neural network;

selecting positive pixel pairs from the patch pair and obtaining positive embedding pairs corresponding to the positive pixel pairs, wherein the positive embedding pairs include global positive embedding pairs and local positive embedding pairs;

for each positive pixel pair:

computing global and local similarity maps using the global and local embedding tensors;

finding global hard negative embeddings using the global similarity maps;

selecting global random negative embeddings from a plurality of patch pairs; and

pooling the global hard negative embeddings and the global random negative embeddings to obtain final global negative embeddings; and

finding local hard negative embeddings using the global and local similarity maps, and randomly sampling final local negative embeddings from the local hard negative embeddings; and

minimizing a final info noise contrastive estimation (InfoNCE) loss including global and local InfoNCE losses, wherein the global InfoNCE loss is computed using the global positive embedding pairs and the final global negative embeddings, and the local InfoNCE loss is computed using the local positive embedding pairs and the final local negative embeddings.

9. The device according to claim 8 , wherein the global and local InfoNCE losses are computed, respectively, according to:

L

=

-

∑

i

=

1

n

p

⁢

o

⁢

s

⁢

log

⁢

exp

⁡

(

f

i

·

f

i

′

/

τ

)

exp

⁡

(

f

i

·

f

i

′

/

τ

)

+

∑

j

=

1

n

n

⁢

e

⁢

g

⁢

exp

⁡

(

f

i

·

h

i

⁢

j

/

τ

)

wherein L denotes a InfoNCE loss, i denotes an i-th element of positive pixel pairs, n pos denotes a number of positive pixels, n neg denotes a number of negative pixels, f i denotes an i-th positive embedding in one of the positive embedding pairs, f i ′ denotes an i-th positive embedding in the other of the positive embedding pairs, τ denotes a temperature parameter, and “.” denotes an inner production operation.

10. The device according to claim 8 , wherein selecting the positive pixel pairs includes:

when one patch of the patch pair overlaps another patch of the patch pair, randomly selecting the positive pixel pairs from an overlapping area of the one patch of the patch pair and the another patch of the patch pair; and

when the one patch of the patch pair does not overlap the another patch of the patch pair, randomly sampling the positive pixel pairs from each patch of the patch pair.

11. The device according to claim 8 , wherein:

for each positive pixel pair, randomly selecting one or more negative pixels from the patch pair, wherein a distance between each of the one or more negative pixels and the each positive pixel pair is greater than a preset value.

12. The device according to claim 9 , wherein the final InfoNCE loss is computed according to:

L final =L g +L l

wherein L final denotes the final InfoNCE loss, L g denotes the global InfoNCE loss, and L l denotes the local InfoNCE loss.

13. A computer program product comprising a non-transitory computer-readable storage medium and program instructions stored therein, the program instructions being configured to be executable by a computer to cause the computer to implement operations comprising:

randomly selecting a plurality of images from an unlabeled image batch;

for each image of the plurality of images,

performing random data augmentation to obtain a patch pair which is inputted to a neural network;

generating global and local embedding tensors for each patch of the patch pair using the neural network;

selecting positive pixel pairs from the patch pair and obtaining positive embedding pairs corresponding to the positive pixel pairs, wherein the positive embedding pairs include global positive embedding pairs and local positive embedding pairs;

for each positive pixel pair:

computing global and local similarity maps using the global and local embedding tensors;

finding global hard negative embeddings using the global similarity maps;

selecting global random negative embeddings from a plurality of patch pairs; and

pooling the global hard negative embeddings and the global random negative embeddings to obtain final global negative embeddings; and

finding local hard negative embeddings using the global and local similarity maps, and randomly sampling final local negative embeddings from the local hard negative embeddings; and

minimizing a final info noise contrastive estimation (InfoNCE) loss including global and local InfoNCE losses, wherein the global InfoNCE loss is computed using the global positive embedding pairs and the final global negative embeddings, and the local InfoNCE loss is computed using the local positive embedding pairs and the final local negative embeddings.

14. The product according to claim 13 , wherein the global and local InfoNCE losses are computed respectively according to:

L

=

-

∑

i

=

1

n

p

⁢

o

⁢

s

⁢

log

⁢

exp

⁡

(

f

i

·

f

i

′

/

τ

)

exp

⁡

(

f

i

·

f

i

′

/

τ

)

+

∑

j

=

1

n

n

⁢

e

⁢

g

⁢

exp

⁡

(

f

i

·

h

i

⁢

j

/

τ

)

wherein L denotes a InfoNCE loss, i denotes an i-th element of positive pixel pairs, n pos denotes a number of positive pixels, n neg denotes a number of negative pixels, f i denotes an i-th positive embedding in one of the positive embedding pairs, f i ′ denotes an i-th positive embedding in the other of the positive embedding pairs, τ denotes a temperature parameter, and “.” denotes an inner production operation.

15. The product according to claim 13 , wherein sampling positive pixel pairs includes:

when one patch of the patch pair overlaps another patch of the patch pair, randomly selecting the positive pixel pairs from an overlapping area of the one patch of the patch pair and the another patch of the patch pair; and

when the one patch of the patch pair does not overlap the another patch of the patch pair, randomly sampling the positive pixel pairs from each patch of the patch pair.

16. The product according to claim 13 , wherein sampling positive pixel pairs includes:

for each positive pixel pair, randomly selecting one or more negative pixels from the patch pair, wherein a distance between each of the one or more negative pixels and the each positive pixel pair is greater than a preset value.

17. The product according to claim 14 , wherein the final InfoNCE loss is computed according to:

L final =L g +L l

wherein L final denotes the final InfoNCE loss, L g denotes the global InfoNCE loss, and L l denotes the local InfoNCE loss.