-- SPDX-License-Identifier: MPL-2.0

-- | 2D resampling/warps, ND binary topology, and exact bounded distance queries.
module HMath.Image.Spatial (
    Interpolation (..),
    resize,
    warp,
    translate,
    rotate,
    Connectivity (..),
    connectedComponents,
    distanceTransform,
) where

import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import HMath.Error
import HMath.Image.Core
import HMath.Image.Filter (validateBoundary)
import qualified HMath.Matrix as M
import qualified HMath.Matrix.LU as LU

{- | Image resampling interpolation mode, choosing nearest,bilinear,bicubic,or Lanczos
kernels.
-}
data Interpolation = Nearest | Bilinear | Bicubic | Lanczos deriving (Interpolation -> Interpolation -> Bool
(Interpolation -> Interpolation -> Bool)
-> (Interpolation -> Interpolation -> Bool) -> Eq Interpolation
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Interpolation -> Interpolation -> Bool
== :: Interpolation -> Interpolation -> Bool
$c/= :: Interpolation -> Interpolation -> Bool
/= :: Interpolation -> Interpolation -> Bool
Eq, Int -> Interpolation -> ShowS
[Interpolation] -> ShowS
Interpolation -> String
(Int -> Interpolation -> ShowS)
-> (Interpolation -> String)
-> ([Interpolation] -> ShowS)
-> Show Interpolation
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Interpolation -> ShowS
showsPrec :: Int -> Interpolation -> ShowS
$cshow :: Interpolation -> String
show :: Interpolation -> String
$cshowList :: [Interpolation] -> ShowS
showList :: [Interpolation] -> ShowS
Show)

{- | 2D resize uses pixel-center mapping (i+0.5)*old/new-0.5, independently
per channel, without antialiasing. Bicubic is Keys a=-0.5; Lanczos has
radius three. Downsampling callers can explicitly prefilter first.
-}
resize :: Interpolation -> Boundary -> (Int, Int) -> Image -> Either MathError Image
resize :: Interpolation
-> Boundary -> (Int, Int) -> Image -> Either MathError Image
resize Interpolation
method Boundary
boundary (Int
height, Int
width) Image
img = do
    Boundary -> Either MathError ()
validateBoundary Boundary
boundary
    Bool -> String -> Either MathError ()
require
        ( [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Image -> [Int]
imageShape Image
img) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2
            Bool -> Bool -> Bool
&& Int
height Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
            Bool -> Bool -> Bool
&& Int
width Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
            Bool -> Bool -> Bool
&& Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
height Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
width Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Image -> Int
imageChannels Image
img) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
262144
        )
        String
"resize shape/storage limit"
    let oldHeight :: Int
oldHeight = Image -> [Int]
imageShape Image
img [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
0; oldWidth :: Int
oldWidth = Image -> [Int]
imageShape Image
img [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
1
    [Int] -> Int -> ([Int] -> Int -> Double) -> Either MathError Image
build
        [Int
height, Int
width]
        (Image -> Int
imageChannels Image
img)
        ( \[Int]
ix Int
c ->
            Interpolation
-> Boundary -> Image -> Double -> Double -> Int -> Double
interpolate
                Interpolation
method
                Boundary
boundary
                Image
img
                ((Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
0) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
oldHeight Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
height Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
0.5)
                ((Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
oldWidth Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
width Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
0.5)
                Int
c
        )

{- | Forward 3x3 homogeneous transform acts on [x(column),y(row),1]. The
matrix is inverted once for output-to-input sampling. Points at projective
infinity or with magnitude over 1e9 are errors, even with constant padding.
-}
warp :: Interpolation -> Boundary -> (Int, Int) -> [[Double]] -> Image -> Either MathError Image
warp :: Interpolation
-> Boundary
-> (Int, Int)
-> [[Double]]
-> Image
-> Either MathError Image
warp Interpolation
method Boundary
boundary (Int
height, Int
width) [[Double]]
transform Image
img = do
    Boundary -> Either MathError ()
validateBoundary Boundary
boundary
    Bool -> String -> Either MathError ()
require
        ( [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Image -> [Int]
imageShape Image
img) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2
            Bool -> Bool -> Bool
&& [[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
transform Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
3
            Bool -> Bool -> Bool
&& ([Double] -> Bool) -> [[Double]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
3) (Int -> Bool) -> ([Double] -> Int) -> [Double] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) [[Double]]
transform
            Bool -> Bool -> Bool
&& Int
height Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
            Bool -> Bool -> Bool
&& Int
width Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
            Bool -> Bool -> Bool
&& Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
height Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
width Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Image -> Int
imageChannels Image
img) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
262144
        )
        String
"warp matrix/output shape"
    _ <- [Double] -> Either MathError [Double]
checked ([[Double]] -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Double]]
transform)
    matrix <- M.fromListChecked 3 3 (concat transform)
    inverse <- M.toList <$> LU.inverse matrix
    points <-
        mapM
            ( \[Int]
ix -> do
                let x :: Double
x = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
1)
                    y :: Double
y = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
0)
                    a :: Double
a = [Double]
inverse [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
0 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double]
inverse [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double]
inverse [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
2
                    b :: Double
b = [Double]
inverse [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
3 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double]
inverse [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
4 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double]
inverse [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
5
                    w :: Double
w = [Double]
inverse [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
6 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double]
inverse [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
7 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double]
inverse [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
8
                _ <- [Double] -> Either MathError [Double]
checked [Double
a, Double
b, Double
w]
                require (w /= 0) "projective point at infinity"
                _ <- checked [a / w, b / w]
                require (abs (a / w) <= 1e9 && abs (b / w) <= 1e9) "warp coordinate limit"
                Right (b / w, a / w)
            )
            (coordinates [height, width])
    image
        [height, width]
        (imageChannels img)
        [interpolate method boundary img y x c | (y, x) <- points, c <- [0 .. imageChannels img - 1]]

-- | Translation in pixel columns/rows; output dimensions equal input.
translate :: Interpolation -> Boundary -> Double -> Double -> Image -> Either MathError Image
translate :: Interpolation
-> Boundary -> Double -> Double -> Image -> Either MathError Image
translate Interpolation
method Boundary
boundary Double
dx Double
dy Image
img = do
    Bool -> String -> Either MathError ()
require ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Image -> [Int]
imageShape Image
img) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2) String
"2D translation required"
    Interpolation
-> Boundary
-> (Int, Int)
-> [[Double]]
-> Image
-> Either MathError Image
warp
        Interpolation
method
        Boundary
boundary
        (Image -> [Int]
imageShape Image
img [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
0, Image -> [Int]
imageShape Image
img [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
1)
        [[Double
1, Double
0, Double
dx], [Double
0, Double
1, Double
dy], [Double
0, Double
0, Double
1]]
        Image
img

{- | Angle in radians, clockwise on row-down image coordinates; rotate about
the center pixel coordinate while preserving the input canvas dimensions.
-}
rotate :: Interpolation -> Boundary -> Double -> Image -> Either MathError Image
rotate :: Interpolation
-> Boundary -> Double -> Image -> Either MathError Image
rotate Interpolation
method Boundary
boundary Double
angle Image
img = do
    Bool -> String -> Either MathError ()
require ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Image -> [Int]
imageShape Image
img) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 Bool -> Bool -> Bool
&& Double -> Bool
finite Double
angle) String
"invalid 2D rotation"
    let cx :: Double
cx = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Image -> [Int]
imageShape Image
img [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2
        cy :: Double
cy = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Image -> [Int]
imageShape Image
img [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2
        a :: Double
a = Double -> Double
forall a. Floating a => a -> a
cos Double
angle
        b :: Double
b = Double -> Double
forall a. Floating a => a -> a
sin Double
angle
    Interpolation
-> Boundary
-> (Int, Int)
-> [[Double]]
-> Image
-> Either MathError Image
warp
        Interpolation
method
        Boundary
boundary
        (Image -> [Int]
imageShape Image
img [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
0, Image -> [Int]
imageShape Image
img [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
1)
        [[Double
a, -Double
b, Double
cx Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
cx Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
cy], [Double
b, Double
a, Double
cy Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
cx Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
cy], [Double
0, Double
0, Double
1]]
        Image
img

interpolate :: Interpolation -> Boundary -> Image -> Double -> Double -> Int -> Double
interpolate :: Interpolation
-> Boundary -> Image -> Double -> Double -> Int -> Double
interpolate Interpolation
Nearest Boundary
boundary Image
img Double
y Double
x Int
c = Boundary -> Image -> [Int] -> Int -> Double
sample Boundary
boundary Image
img [Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5), Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5)] Int
c
interpolate Interpolation
method Boundary
boundary Image
img Double
y Double
x Int
c =
    let radius :: Int
radius = case Interpolation
method of Interpolation
Bilinear -> Int
1; Interpolation
Bicubic -> Int
2; Interpolation
Lanczos -> Int
3
        weights :: b -> [(Int, b)]
weights b
value = [(Int
i, b -> b
forall {a}. (Ord a, Floating a) => a -> a
kernel (b
value b -> b -> b
forall a. Num a => a -> a -> a
- Int -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i)) | Int
i <- [b -> Int
forall b. Integral b => b -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor b
value Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
radius Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 .. b -> Int
forall b. Integral b => b -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor b
value Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
radius]]
        kernel :: a -> a
kernel a
t =
            let a :: a
a = a -> a
forall a. Num a => a -> a
abs a
t
             in case Interpolation
method of
                    Interpolation
Bilinear -> a -> a -> a
forall a. Ord a => a -> a -> a
max a
0 (a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
a)
                    Interpolation
Bicubic ->
                        if a
a a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
1
                            then a
1.5 a -> a -> a
forall a. Num a => a -> a -> a
* a
a a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
3 :: Int) a -> a -> a
forall a. Num a => a -> a -> a
- a
2.5 a -> a -> a
forall a. Num a => a -> a -> a
* a
a a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) a -> a -> a
forall a. Num a => a -> a -> a
+ a
1
                            else if a
a a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
2 then (-a
0.5) a -> a -> a
forall a. Num a => a -> a -> a
* a
a a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
3 :: Int) a -> a -> a
forall a. Num a => a -> a -> a
+ a
2.5 a -> a -> a
forall a. Num a => a -> a -> a
* a
a a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) a -> a -> a
forall a. Num a => a -> a -> a
- a
4 a -> a -> a
forall a. Num a => a -> a -> a
* a
a a -> a -> a
forall a. Num a => a -> a -> a
+ a
2 else a
0
                    Interpolation
Lanczos ->
                        if a
a a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 then a
1 else if a
a a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
3 then a -> a
forall a. Floating a => a -> a
sin (a
forall a. Floating a => a
pi a -> a -> a
forall a. Num a => a -> a -> a
* a
a) a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
forall a. Floating a => a
pi a -> a -> a
forall a. Num a => a -> a -> a
* a
a) a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
sin (a
forall a. Floating a => a
pi a -> a -> a
forall a. Num a => a -> a -> a
* a
a a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
3) a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
forall a. Floating a => a
pi a -> a -> a
forall a. Num a => a -> a -> a
* a
a a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
3) else a
0
        wy :: [(Int, Double)]
wy = Double -> [(Int, Double)]
forall {b}. (RealFrac b, Floating b) => b -> [(Int, b)]
weights Double
y
        wx :: [(Int, Double)]
wx = Double -> [(Int, Double)]
forall {b}. (RealFrac b, Floating b) => b -> [(Int, b)]
weights Double
x
        total :: Double
total = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (((Int, Double) -> Double) -> [(Int, Double)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Double) -> Double
forall a b. (a, b) -> b
snd [(Int, Double)]
wy) Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (((Int, Double) -> Double) -> [(Int, Double)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Double) -> Double
forall a b. (a, b) -> b
snd [(Int, Double)]
wx)
     in [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
total Double -> Double -> Double
forall a. Num a => a -> a -> a
* Boundary -> Image -> [Int] -> Int -> Double
sample Boundary
boundary Image
img [Int
i, Int
j] Int
c | (Int
i, Double
a) <- [(Int, Double)]
wy, (Int
j, Double
b) <- [(Int, Double)]
wx]

-- | Spatial connected-component adjacency: axis neighbors or all surrounding neighbors.
data Connectivity = AxisConnected | FullyConnected deriving (Connectivity -> Connectivity -> Bool
(Connectivity -> Connectivity -> Bool)
-> (Connectivity -> Connectivity -> Bool) -> Eq Connectivity
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Connectivity -> Connectivity -> Bool
== :: Connectivity -> Connectivity -> Bool
$c/= :: Connectivity -> Connectivity -> Bool
/= :: Connectivity -> Connectivity -> Bool
Eq, Int -> Connectivity -> ShowS
[Connectivity] -> ShowS
Connectivity -> String
(Int -> Connectivity -> ShowS)
-> (Connectivity -> String)
-> ([Connectivity] -> ShowS)
-> Show Connectivity
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Connectivity -> ShowS
showsPrec :: Int -> Connectivity -> ShowS
$cshow :: Connectivity -> String
show :: Connectivity -> String
$cshowList :: [Connectivity] -> ShowS
showList :: [Connectivity] -> ShowS
Show)

{- | Nonzero is foreground. Label zero is background; component labels start
at one in first-pixel row-major order. Returns labels as exact Doubles and
component sizes in that same order. Connectivity has no wraparound.
-}
connectedComponents :: Connectivity -> Image -> Either MathError (Image, [Int])
connectedComponents :: Connectivity -> Image -> Either MathError (Image, [Int])
connectedComponents Connectivity
connectivity Image
img = do
    Bool -> String -> Either MathError ()
require (Image -> Int
imageChannels Image
img Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1) String
"components need one channel"
    let dims :: [Int]
dims = Image -> [Int]
imageShape Image
img; rank :: Int
rank = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims
    Bool -> String -> Either MathError ()
require (Int
rank Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
4) String
"component rank limit (four)"
    let offsets :: [[Int]]
offsets = case Connectivity
connectivity of
            Connectivity
AxisConnected ->
                [Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
k (Int -> [Int]
forall a. a -> [a]
repeat Int
0) [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
sign] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate (Int
rank Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
0 | Int
k <- [Int
0 .. Int
rank Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Int
sign <- [-Int
1, Int
1]]
            Connectivity
FullyConnected -> [(Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1) [Int]
ix | [Int]
ix <- [Int] -> [[Int]]
coordinates (Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate Int
rank Int
3), (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
1) [Int]
ix]
        foreground :: Set [Int]
foreground = [[Int]] -> Set [Int]
forall a. Ord a => [a] -> Set a
Set.fromList [[Int]
ix | [Int]
ix <- [Int] -> [[Int]]
coordinates [Int]
dims, Image -> [Int] -> Int -> Double
at Image
img [Int]
ix Int
0 Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
/= Double
0]
        flood :: [[Int]] -> Set [Int] -> [[Int]] -> (Set [Int], [[Int]])
flood [] Set [Int]
pending [[Int]]
found = (Set [Int]
pending, [[Int]]
found)
        flood ([Int]
ix : [[Int]]
queue) Set [Int]
pending [[Int]]
found =
            let next :: [[Int]]
next = ([Int] -> Bool) -> [[Int]] -> [[Int]]
forall a. (a -> Bool) -> [a] -> [a]
filter ([Int] -> Set [Int] -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set [Int]
pending) [(Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [Int]
ix [Int]
offset | [Int]
offset <- [[Int]]
offsets]
                remaining :: Set [Int]
remaining = ([Int] -> Set [Int] -> Set [Int])
-> Set [Int] -> [[Int]] -> Set [Int]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr [Int] -> Set [Int] -> Set [Int]
forall a. Ord a => a -> Set a -> Set a
Set.delete Set [Int]
pending [[Int]]
next
             in [[Int]] -> Set [Int] -> [[Int]] -> (Set [Int], [[Int]])
flood ([[Int]]
next [[Int]] -> [[Int]] -> [[Int]]
forall a. [a] -> [a] -> [a]
++ [[Int]]
queue) Set [Int]
remaining ([Int]
ix [Int] -> [[Int]] -> [[Int]]
forall a. a -> [a] -> [a]
: [[Int]]
found)
        components :: Set [Int] -> Map [Int] a -> [Int] -> t -> (Map [Int] a, [Int])
components Set [Int]
pending Map [Int] a
labels [Int]
sizes t
label
            | Set [Int] -> Bool
forall a. Set a -> Bool
Set.null Set [Int]
pending = (Map [Int] a
labels, [Int] -> [Int]
forall a. [a] -> [a]
reverse [Int]
sizes)
            | Bool
otherwise =
                let start :: [Int]
start = Set [Int] -> [Int]
forall a. Set a -> a
Set.findMin Set [Int]
pending
                    (Set [Int]
remaining, [[Int]]
found) = [[Int]] -> Set [Int] -> [[Int]] -> (Set [Int], [[Int]])
flood [[Int]
start] ([Int] -> Set [Int] -> Set [Int]
forall a. Ord a => a -> Set a -> Set a
Set.delete [Int]
start Set [Int]
pending) []
                    next :: Map [Int] a
next = ([Int] -> Map [Int] a -> Map [Int] a)
-> Map [Int] a -> [[Int]] -> Map [Int] a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\[Int]
ix -> [Int] -> a -> Map [Int] a -> Map [Int] a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert [Int]
ix (t -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral t
label)) Map [Int] a
labels [[Int]]
found
                 in Set [Int] -> Map [Int] a -> [Int] -> t -> (Map [Int] a, [Int])
components Set [Int]
remaining Map [Int] a
next ([[Int]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Int]]
found Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
sizes) (t
label t -> t -> t
forall a. Num a => a -> a -> a
+ t
1)
        (Map [Int] Double
finalLabels, [Int]
finalSizes) = Set [Int]
-> Map [Int] Double -> [Int] -> Int -> (Map [Int] Double, [Int])
forall {t} {a}.
(Integral t, Num a) =>
Set [Int] -> Map [Int] a -> [Int] -> t -> (Map [Int] a, [Int])
components Set [Int]
foreground Map [Int] Double
forall k a. Map k a
Map.empty [] (Int
1 :: Int)
    Bool -> String -> Either MathError ()
require
        (Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Set [Int] -> Int
forall a. Set a -> Int
Set.size Set [Int]
foreground) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([[Int]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Int]]
offsets) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
20000000)
        String
"component work limit"
    result <- [Int] -> Int -> ([Int] -> Int -> Double) -> Either MathError Image
build [Int]
dims Int
1 (\[Int]
ix Int
_ -> Double -> [Int] -> Map [Int] Double -> Double
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Double
0 [Int]
ix Map [Int] Double
finalLabels)
    Right (result, finalSizes)

{- | Exact brute-force Euclidean distance to the nearest zero-valued spatial
sample, with strictly positive per-axis spacing. No implicit outside zeros;
an all-foreground image has no finite answer and returns DomainError.
-}
distanceTransform :: [Double] -> Image -> Either MathError Image
distanceTransform :: [Double] -> Image -> Either MathError Image
distanceTransform [Double]
spacing Image
img = do
    Bool -> String -> Either MathError ()
require
        ( Image -> Int
imageChannels Image
img Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
            Bool -> Bool -> Bool
&& [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
spacing Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Image -> [Int]
imageShape Image
img)
            Bool -> Bool -> Bool
&& (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Double
x -> Double -> Bool
finite Double
x Bool -> Bool -> Bool
&& Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0) [Double]
spacing
            Bool -> Bool -> Bool
&& Image -> Int
spatialSize Image
img Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
4096
        )
        String
"distance transform shape/spacing/size"
    let zeros :: [[Int]]
zeros = [[Int]
ix | [Int]
ix <- [Int] -> [[Int]]
coordinates (Image -> [Int]
imageShape Image
img), Image -> [Int] -> Int -> Double
at Image
img [Int]
ix Int
0 Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0]
    Bool -> String -> Either MathError ()
require (Bool -> Bool
not ([[Int]] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Int]]
zeros)) String
"distance transform has no background sample"
    let norm :: [t] -> t
norm [t]
values =
            let scale :: t
scale = [t] -> t
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (t
0 t -> [t] -> [t]
forall a. a -> [a] -> [a]
: (t -> t) -> [t] -> [t]
forall a b. (a -> b) -> [a] -> [b]
map t -> t
forall a. Num a => a -> a
abs [t]
values)
             in if t
scale t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
0 then t
0 else t
scale t -> t -> t
forall a. Num a => a -> a -> a
* t -> t
forall a. Floating a => a -> a
sqrt ([t] -> t
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [(t
x t -> t -> t
forall a. Fractional a => a -> a -> a
/ t
scale) t -> Int -> t
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) | t
x <- [t]
values])
    [Int] -> Int -> ([Int] -> Int -> Double) -> Either MathError Image
build
        (Image -> [Int]
imageShape Image
img)
        Int
1
        ( \[Int]
ix Int
_ ->
            [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [[Double] -> Double
forall {t}. (Floating t, Ord t) => [t] -> t
norm ((Int -> Int -> Double -> Double)
-> [Int] -> [Int] -> [Double] -> [Double]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\Int
a Int
b Double
step -> Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
a Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
b) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
step) [Int]
ix [Int]
zero [Double]
spacing) | [Int]
zero <- [[Int]]
zeros]
        )