from ddd_pylib._base import ImageOperator
import xarray as xr
import numpy as np
from typing import Callable
[docs]
class ImageCalculatorOperator(ImageOperator):
"""
Applies a binary operation between two images.
The operation can be one of the following: Add, Subtract, Multiply, Divide, OR, AND, Difference, Minimum, Maximum, or Average.
The data type of the output image is the one used by the main image (left hand of the operation).
All operations are clipped to the data type for integer types.
The AND and OR operations are numeric. They work as expected for booleans, but for numeric values,
they return the first non-zero value (for OR) or the second value if the first is non-zero (for AND).
The division is not secured against division by zero.
"""
operations = set(
[
"Add",
"Subtract",
"Multiply",
"Divide",
"OR",
"AND",
"Difference",
"Minimum",
"Maximum",
"Average",
]
)
def __init__(self):
super().__init__()
self._secondaryImage = None
self._currentOperationName = ""
self._currentOperation = None
self._availableOperations = {
"Add": self._add,
"Subtract": self._subtract,
"Multiply": self._multiply,
"Divide": self._divide,
"OR": self._numericOr,
"AND": self._numericAnd,
"Difference": self._subtractAbs,
"Minimum": self._minimum,
"Maximum": self._maximum,
"Average": self._average,
}
self.setOperation(self.defaultOperationName())
[docs]
@staticmethod
def defaultOperationName() -> str:
"""
Set the Add operation by default.
"""
return "Add"
[docs]
def setSecondaryImage(self, img: xr.DataArray):
"""
Set the secondary image for the binary operation.
Args:
img: The secondary image
"""
self._secondaryImage = img
[docs]
def setOperation(self, operationName: str):
"""
Set the operation that needs to be executed
Args:
operationName: name of the operation to execute
"""
if operationName not in self._availableOperations:
raise ValueError("Operation not valid")
self._currentOperationName = operationName
self._currentOperation = self._availableOperations[operationName]
[docs]
def getSecondaryImage(self) -> xr.DataArray:
"""
Return the secondary image used in the operation
"""
if self._secondaryImage is None:
raise ValueError("Secondary image is not set")
return self._secondaryImage
[docs]
def getPrefix(self) -> str:
return f"Calc[{self._currentOperationName}]-"
[docs]
def getCurrentOperation(
self,
) -> Callable[[xr.DataArray, xr.DataArray], xr.DataArray]:
"""
Return the current operation function
"""
if self._currentOperation is None:
raise ValueError("No operation has been set.")
return self._currentOperation
[docs]
def getCurrentOperationName(self) -> str:
"""
Return the name of the current operation
"""
return self._currentOperationName
[docs]
def sanityCheck(self):
super().sanityCheck()
imA = self.getMainImage()
imB = self.getSecondaryImage()
if imB is None:
raise ValueError("The secondary image is missing.")
if (
self._currentOperationName not in self._availableOperations
or self._currentOperation is None
):
raise ValueError("No operation has been set.")
if imA.sizes != imB.sizes:
raise ValueError("Both images must have the same size.")
def _add(self, imA: xr.DataArray, imB: xr.DataArray) -> xr.DataArray:
return xr.ufuncs.add(imA, imB)
def _subtract(self, imA: xr.DataArray, imB: xr.DataArray) -> xr.DataArray:
return xr.ufuncs.subtract(imA, imB)
def _multiply(self, imA: xr.DataArray, imB: xr.DataArray) -> xr.DataArray:
return xr.ufuncs.multiply(imA, imB)
def _divide(self, imA: xr.DataArray, imB: xr.DataArray) -> xr.DataArray:
return xr.ufuncs.divide(imA, imB)
def _numericOr(self, imA: xr.DataArray, imB: xr.DataArray) -> xr.DataArray:
a = imA.values
b = imB.values
r = np.where(~np.isclose(a, 0, atol=1e-6), a, b)
return xr.DataArray(r, dims=imA.dims, attrs=imA.attrs)
def _numericAnd(self, imA: xr.DataArray, imB: xr.DataArray) -> xr.DataArray:
a = imA.values
b = imB.values
r = np.where(~np.isclose(a, 0, atol=1e-6), b, a)
return xr.DataArray(r, dims=imA.dims, attrs=imA.attrs)
def _subtractAbs(self, imA: xr.DataArray, imB: xr.DataArray) -> xr.DataArray:
return xr.ufuncs.abs(self._subtract(imA, imB))
def _minimum(self, imA: xr.DataArray, imB: xr.DataArray) -> xr.DataArray:
return xr.ufuncs.minimum(imA, imB)
def _maximum(self, imA: xr.DataArray, imB: xr.DataArray) -> xr.DataArray:
return xr.ufuncs.maximum(imA, imB)
def _average(self, imA: xr.DataArray, imB: xr.DataArray) -> xr.DataArray:
return (imA + imB) / 2
def _clippedOperation(self, imA: xr.DataArray, imB: xr.DataArray) -> xr.DataArray:
ori_type = imA.dtype
imA = imA.astype(np.float32)
imB = imB.astype(np.float32)
operation = self.getCurrentOperation()
result = operation(imA, imB)
if np.issubdtype(ori_type, np.integer):
info = np.iinfo(ori_type)
result.values = np.clip(result.values, info.min, info.max)
return result.astype(ori_type)
def _applyToFrame(self, frameData, t, nT):
imA, imB = frameData
if self._currentOperation is None:
raise ValueError("No operation has been set.")
return self._clippedOperation(imA, imB.transpose(*imA.dims))
def _zipInputs(self):
imA = self.getMainImage()
imB = self.getSecondaryImage()
if "T" in imA.dims:
return [
(imA.isel(T=t), imB.isel(T=t)) for t in range(self.getNTimePoints())
]
else:
return [(imA, imB)]