mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-10 21:42:39 +02:00
61 lines
1.9 KiB
Python
61 lines
1.9 KiB
Python
from __future__ import division, print_function, absolute_import
|
|
import numpy as np
|
|
import torch as t
|
|
from CDTools.tools import cmath
|
|
|
|
|
|
def centroid(im, dims=2):
|
|
"""Returns the centroid of an image or a stack of images
|
|
|
|
By default, the last two dimensions are used in the calculation
|
|
and the remainder of the dimensions are passed through.
|
|
|
|
Beware that the meaning of the centroid is not well defined if your
|
|
image contains values less than 0
|
|
|
|
Args:
|
|
im (t.Tensor) : An image or stack of images to calculate from
|
|
dims (int) : Default 2, how many trailing dimensions to calculate for
|
|
|
|
Returns:
|
|
t.Tensor : An (i,j) index or stack of indices
|
|
"""
|
|
indices = (t.arange(im.shape[-dims+i]).to(t.float32) for i in range(dims))
|
|
indices = t.meshgrid(*indices)
|
|
|
|
use_dims = [-dims+i for i in range(dims)]
|
|
divisor = t.sum(im, dim=use_dims)
|
|
centroids = [t.sum(index * im, dim=use_dims) / divisor
|
|
for index in indices]
|
|
|
|
return t.stack(centroids,dim=-1)
|
|
|
|
|
|
|
|
|
|
def centroid_sq(im, dims=2, comp=False):
|
|
"""Returns the centroid of the square of an image or stack of images
|
|
|
|
By default, the last two dimensions are used in the calculation
|
|
and the remainder of the dimensions are passed through.
|
|
|
|
If the "comp" flag is set, it will be assumed that the last dimension
|
|
represents the real and imaginary part of a complex number, and the
|
|
centroid will be calculated for the magnitude squared of those numbers
|
|
|
|
Args:
|
|
im (t.Tensor) : An image or stack of images to calculate from
|
|
dims (int) : Default 2, how many trailing dimensions to calculate for
|
|
comp (bool) : Default is False, whether the data represents complex numbers
|
|
Returns:
|
|
t.Tensor : An (i,j) index or stack of indices
|
|
"""
|
|
if comp:
|
|
im_sq = cmath.cabssq(im)
|
|
else:
|
|
im_sq = im**2
|
|
|
|
return centroid(im_sq, dims=dims)
|
|
|
|
|