{-# LANGUAGE BangPatterns
           , FlexibleContexts
           , MultiWayIf #-}

module Vision.Detector.Edge (canny) where

import Control.Monad (when)
import Control.Monad.ST (ST)
import Data.Int
import Data.Vector.Storable (enumFromN, forM_)
import Foreign.Storable (Storable)

import Vision.Image (
      Image, ImagePixel, Manifest, MutableManifest, Grey, DerivativeType (..)
    , (!), shape, linearIndex, fromFunction
    , create, new', linearRead, linearWrite
    , sobel
    )
import Vision.Primitive (Z (..), (:.) (..), inShape, ix2)

data EdgeDirection = NorthSouth         -- ^ |
                   | WestEast           -- ^ ―
                   | NorthEastSouthWest -- ^ /
                   | NorthWestSouthEast -- ^ \

-- | Detects edges using the Canny's algorithm. Edges are given the value
-- 'maxBound' while non-edges are given the value 'minBound'.
--
-- This implementation doesn't perform any noise erasing (as blurring) before
-- edge detection. Noisy images might need to be pre-processed using a Gaussian
-- blur.
--
-- The bidirectional derivative (gradient magnitude) is computed from @x@ and
-- @y@ derivatives using @sqrt(dx² + dy²)@.
--
-- See <http://en.wikipedia.org/wiki/Canny_edge_detector> for details.
--
-- This function is specialized for 'Grey' images but is declared @INLINABLE@
-- to be further specialized for new image types.
canny :: ( Image src, Integral (ImagePixel src), Bounded res, Eq res
         , Storable res)
      => Int
      -- ^ Radius of the Sobel's filter.
      -> Int32
      -- ^ Low threshold. Pixels for which the bidirectional derivative is
      -- greater than this value and which are connected to another pixel which
      -- is part of an edge will be part of this edge.
      -> Int32
      -- ^ High threshold. Pixels for which the bidirectional derivative is
      -- greater than this value will be part of an edge.
      -> src
      -> Manifest res
canny :: forall src res.
(Image src, Integral (ImagePixel src), Bounded res, Eq res,
 Storable res) =>
Int -> Int32 -> Int32 -> src -> Manifest res
canny !Int
derivSize !Int32
lowThres !Int32
highThres !src
img =
    forall (i :: * -> *).
MutableImage i =>
(forall s. ST s (i s)) -> Freezed i
create forall a b. (a -> b) -> a -> b
$ do
        MutableManifest res s
edges <- forall p s. (Storable p, Bounded p) => ST s (MutableManifest p s)
newManifest
        forall (m :: * -> *) a b.
(Monad m, Storable a) =>
Vector a -> (a -> m b) -> m ()
forM_ (forall a. (Storable a, Num a) => a -> Int -> Vector a
enumFromN Int
0 Int
h) forall a b. (a -> b) -> a -> b
$ \Int
y -> do
            let !lineOffset :: Int
lineOffset = Int
y forall a. Num a => a -> a -> a
* Int
w
            forall (m :: * -> *) a b.
(Monad m, Storable a) =>
Vector a -> (a -> m b) -> m ()
forM_ (forall a. (Storable a, Num a) => a -> Int -> Vector a
enumFromN Int
0 Int
w) forall a b. (a -> b) -> a -> b
$ \Int
x -> do
                forall {m :: * -> *} {i :: * -> *}.
(MutableImage i, PrimMonad m, Eq (ImagePixel (Freezed i)),
 Bounded (ImagePixel (Freezed i))) =>
i (PrimState m) -> Int -> Int -> Int -> Int32 -> m ()
visitPoint MutableManifest res s
edges Int
x Int
y (Int
lineOffset forall a. Num a => a -> a -> a
+ Int
x) Int32
highThres'
        forall (m :: * -> *) a. Monad m => a -> m a
return MutableManifest res s
edges
  where
    !size :: Size
size@(DIM0
Z :. Int
h :. Int
w) = forall i. MaskedImage i => i -> Size
shape src
img

    -- Squares both thresholds as they will be compared to 'dxy' which contains
    -- squared gradient magnitudes.
    (!Int32
lowThres', !Int32
highThres') = (forall a. Num a => a -> a
square Int32
lowThres, forall a. Num a => a -> a
square Int32
highThres)

    dx, dy :: Manifest Int16
    !dx :: Manifest Int16
dx = forall src res.
(Image src, Integral (ImagePixel src), FromFunction res,
 Integral (FromFunctionPixel res), Storable (FromFunctionPixel res),
 SeparatelyFiltrable src res (FromFunctionPixel res)) =>
Int -> DerivativeType -> src -> res
sobel Int
derivSize DerivativeType
DerivativeX src
img
    !dy :: Manifest Int16
dy = forall src res.
(Image src, Integral (ImagePixel src), FromFunction res,
 Integral (FromFunctionPixel res), Storable (FromFunctionPixel res),
 SeparatelyFiltrable src res (FromFunctionPixel res)) =>
Int -> DerivativeType -> src -> res
sobel Int
derivSize DerivativeType
DerivativeY src
img

    -- Gradient magnitude, squared.
    dxy :: Manifest Int32
    !dxy :: Manifest Int32
dxy = forall i.
FromFunction i =>
Size -> (Size -> FromFunctionPixel i) -> i
fromFunction Size
size forall a b. (a -> b) -> a -> b
$ \Size
pt ->
                  forall a. Num a => a -> a
square (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ Manifest Int16
dx forall i. Image i => i -> Size -> ImagePixel i
! Size
pt)
                forall a. Num a => a -> a -> a
+ forall a. Num a => a -> a
square (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ Manifest Int16
dy forall i. Image i => i -> Size -> ImagePixel i
! Size
pt)

    newManifest :: (Storable p, Bounded p) => ST s (MutableManifest p s)
    newManifest :: forall p s. (Storable p, Bounded p) => ST s (MutableManifest p s)
newManifest = forall (i :: * -> *) (m :: * -> *).
(MutableImage i, PrimMonad m) =>
Size -> ImagePixel (Freezed i) -> m (i (PrimState m))
new' Size
size forall a. Bounded a => a
minBound

    -- Visits a point and compares its gradient magnitude to the given
    -- threshold, visits neighbor if the point is perceived an an edge.
    visitPoint :: i (PrimState m) -> Int -> Int -> Int -> Int32 -> m ()
visitPoint !i (PrimState m)
edges !Int
x !Int
y !Int
linearIX !Int32
thres = do
        ImagePixel (Freezed i)
val <- forall (i :: * -> *) (m :: * -> *).
(MutableImage i, PrimMonad m) =>
i (PrimState m) -> Int -> m (ImagePixel (Freezed i))
linearRead i (PrimState m)
edges Int
linearIX

        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ImagePixel (Freezed i)
val forall a. Eq a => a -> a -> Bool
== forall a. Bounded a => a
minBound) forall a b. (a -> b) -> a -> b
$ do
            let !ptDxy :: ImagePixel (Manifest Int32)
ptDxy    = Manifest Int32
dxy forall i. Image i => i -> Int -> ImagePixel i
`linearIndex` Int
linearIX
                ptDx :: ImagePixel (Manifest Int16)
ptDx      = Manifest Int16
dx  forall i. Image i => i -> Int -> ImagePixel i
`linearIndex` Int
linearIX
                ptDy :: ImagePixel (Manifest Int16)
ptDy      = Manifest Int16
dy  forall i. Image i => i -> Int -> ImagePixel i
`linearIndex` Int
linearIX
                direction :: EdgeDirection
direction = forall {p} {p}. (Integral p, Integral p) => p -> p -> EdgeDirection
edgeDirection Int16
ptDx Int16
ptDy

            -- When the current pixel has a greater magnitude than the threshold
            -- and is a local maximum, considers it as a new starting point of
            -- an edge. Tries to draw the remaining of the edge using the low
            -- threshold and by following the edge direction.

            forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int32
ptDxy forall a. Ord a => a -> a -> Bool
> Int32
thres Bool -> Bool -> Bool
&& forall {t}.
(Num t, Ord t) =>
Int -> Int -> t -> EdgeDirection -> Bool
isMaximum Int
x Int
y Int32
ptDxy EdgeDirection
direction) forall a b. (a -> b) -> a -> b
$ do
                forall (i :: * -> *) (m :: * -> *).
(MutableImage i, PrimMonad m) =>
i (PrimState m) -> Int -> ImagePixel (Freezed i) -> m ()
linearWrite i (PrimState m)
edges Int
linearIX forall a. Bounded a => a
maxBound
                i (PrimState m) -> Int -> Int -> EdgeDirection -> m ()
visitNeighbour i (PrimState m)
edges Int
x Int
y EdgeDirection
direction

    visitNeighbour :: i (PrimState m) -> Int -> Int -> EdgeDirection -> m ()
visitNeighbour !i (PrimState m)
edges !Int
x !Int
y !EdgeDirection
direction = do
        let (!Int
x1, !Int
y1, !Int
x2, !Int
y2) =
                case EdgeDirection
direction of
                    EdgeDirection
NorthSouth         -> (Int
x,     Int
y forall a. Num a => a -> a -> a
- Int
1, Int
x,     Int
y forall a. Num a => a -> a -> a
+ Int
1)
                    EdgeDirection
WestEast           -> (Int
x forall a. Num a => a -> a -> a
- Int
1, Int
y,     Int
x forall a. Num a => a -> a -> a
+ Int
1, Int
y    )
                    EdgeDirection
NorthEastSouthWest -> (Int
x forall a. Num a => a -> a -> a
- Int
1, Int
y forall a. Num a => a -> a -> a
- Int
1, Int
x forall a. Num a => a -> a -> a
+ Int
1, Int
y forall a. Num a => a -> a -> a
+ Int
1)
                    EdgeDirection
NorthWestSouthEast -> (Int
x forall a. Num a => a -> a -> a
+ Int
1, Int
y forall a. Num a => a -> a -> a
- Int
1, Int
x forall a. Num a => a -> a -> a
- Int
1, Int
y forall a. Num a => a -> a -> a
+ Int
1)

        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall sh. Shape sh => sh -> sh -> Bool
inShape Size
size (Int -> Int -> Size
ix2 Int
y1 Int
x1)) forall a b. (a -> b) -> a -> b
$
            i (PrimState m) -> Int -> Int -> Int -> Int32 -> m ()
visitPoint i (PrimState m)
edges Int
x1 Int
y1 (Int
y1 forall a. Num a => a -> a -> a
* Int
w forall a. Num a => a -> a -> a
+ Int
x1) Int32
lowThres'

        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall sh. Shape sh => sh -> sh -> Bool
inShape Size
size (Int -> Int -> Size
ix2 Int
y2 Int
x2)) forall a b. (a -> b) -> a -> b
$
            i (PrimState m) -> Int -> Int -> Int -> Int32 -> m ()
visitPoint i (PrimState m)
edges Int
x2 Int
y2 (Int
y2 forall a. Num a => a -> a -> a
* Int
w forall a. Num a => a -> a -> a
+ Int
x2) Int32
lowThres'

    isMaximum :: Int -> Int -> t -> EdgeDirection -> Bool
isMaximum !Int
x !Int
y !t
ptDxy !EdgeDirection
direction =
        let (!Int
x1, !Int
y1, !Int
x2, !Int
y2) =
                case EdgeDirection
direction of
                    EdgeDirection
NorthSouth         -> (Int
x forall a. Num a => a -> a -> a
- Int
1, Int
y,     Int
x forall a. Num a => a -> a -> a
+ Int
1, Int
y    )
                    EdgeDirection
WestEast           -> (Int
x,     Int
y forall a. Num a => a -> a -> a
- Int
1, Int
x,     Int
y forall a. Num a => a -> a -> a
+ Int
1)
                    EdgeDirection
NorthEastSouthWest -> (Int
x forall a. Num a => a -> a -> a
+ Int
1, Int
y forall a. Num a => a -> a -> a
- Int
1, Int
x forall a. Num a => a -> a -> a
- Int
1, Int
y forall a. Num a => a -> a -> a
+ Int
1)
                    EdgeDirection
NorthWestSouthEast -> (Int
x forall a. Num a => a -> a -> a
- Int
1, Int
y forall a. Num a => a -> a -> a
- Int
1, Int
x forall a. Num a => a -> a -> a
+ Int
1, Int
y forall a. Num a => a -> a -> a
+ Int
1)
        in forall {t} {t}.
Num t =>
t -> (t -> t -> Bool) -> (Int, Int) -> Bool
tryCompare t
ptDxy forall a. Ord a => a -> a -> Bool
(>) (Int
x1, Int
y1) Bool -> Bool -> Bool
&& forall {t} {t}.
Num t =>
t -> (t -> t -> Bool) -> (Int, Int) -> Bool
tryCompare t
ptDxy forall a. Ord a => a -> a -> Bool
(>=) (Int
x2, Int
y2)

    tryCompare :: t -> (t -> t -> Bool) -> (Int, Int) -> Bool
tryCompare !t
ptDxy t -> t -> Bool
op !(Int
x, Int
y)
        | forall sh. Shape sh => sh -> sh -> Bool
inShape Size
size (Int -> Int -> Size
ix2 Int
y Int
x) = t
ptDxy t -> t -> Bool
`op` forall a b. (Integral a, Num b) => a -> b
fromIntegral (Manifest Int32
dxy forall i. Image i => i -> Size -> ImagePixel i
! Int -> Int -> Size
ix2 Int
y Int
x)
        | Bool
otherwise              = Bool
True

    -- Returns the direction of the edge, not to be confused with the direction
    -- of the gradient which is the perpendicular of this value.
    edgeDirection :: p -> p -> EdgeDirection
edgeDirection p
ptDx p
ptDy =
        let !angle :: Double
angle = forall a. RealFloat a => a -> a -> a
atan2 (forall a. Integral a => a -> Double
double p
ptDy) (forall a. Integral a => a -> Double
double p
ptDx)
        in if Double
angle forall a. Ord a => a -> a -> Bool
>= Double
0 then if | Double
angle forall a. Ord a => a -> a -> Bool
>  Double
pi8x7 -> EdgeDirection
NorthSouth
                                 | Double
angle forall a. Ord a => a -> a -> Bool
>  Double
pi8x5 -> EdgeDirection
NorthEastSouthWest
                                 | Double
angle forall a. Ord a => a -> a -> Bool
>  Double
pi8x3 -> EdgeDirection
WestEast
                                 | Double
angle forall a. Ord a => a -> a -> Bool
>    Double
pi8 -> EdgeDirection
NorthWestSouthEast
                                 | Bool
otherwise      -> EdgeDirection
NorthSouth
                         else if | Double
angle forall a. Ord a => a -> a -> Bool
< -Double
pi8x7 -> EdgeDirection
NorthSouth
                                 | Double
angle forall a. Ord a => a -> a -> Bool
< -Double
pi8x5 -> EdgeDirection
NorthWestSouthEast
                                 | Double
angle forall a. Ord a => a -> a -> Bool
< -Double
pi8x3 -> EdgeDirection
WestEast
                                 | Double
angle forall a. Ord a => a -> a -> Bool
<   -Double
pi8 -> EdgeDirection
NorthEastSouthWest
                                 | Bool
otherwise      -> EdgeDirection
NorthSouth

    !pi8 :: Double
pi8   = forall a. Floating a => a
pi forall a. Fractional a => a -> a -> a
/ Double
8
    !pi8x3 :: Double
pi8x3 = Double
pi8 forall a. Num a => a -> a -> a
* Double
3
    !pi8x5 :: Double
pi8x5 = Double
pi8 forall a. Num a => a -> a -> a
* Double
5
    !pi8x7 :: Double
pi8x7 = Double
pi8 forall a. Num a => a -> a -> a
* Double
7
{-# INLINABLE  canny #-}
{-# SPECIALIZE canny :: Int -> Int32 -> Int32 -> Grey -> Grey #-}

square :: Num a => a -> a
square :: forall a. Num a => a -> a
square a
a = a
a forall a. Num a => a -> a -> a
* a
a

double :: Integral a => a -> Double
double :: forall a. Integral a => a -> Double
double = forall a b. (Integral a, Num b) => a -> b
fromIntegral