content: {-# LANGUAGE ConstraintKinds   #-}
  {-# LANGUAGE FlexibleContexts  #-}
  {-# LANGUAGE ViewPatterns      #-}
  -- An implementation of the Canny edge detection algorithm
  --
  --   J. F. Canny, "A Computational Approach to Edge Detection" in _Pattern
  --   Analysis and Machine Intelligence_, (6), 1986.
  --
  -- This module defines the first seven data-parallel sections of the algorithm.
  -- The last phase, which uses a recursive algorithm to "connect" pixels that
  -- form the output image, is defined in the module Wildfire and implemented
  -- using Repa.
  --
  
  module Canny where
  
  import qualified Prelude                                as P
  
  import Data.Array.Accelerate                            as A
  import Data.Array.Accelerate.IO.Codec.BMP               as A
  import Data.Array.Accelerate.Data.Colour.RGB
  
  
  -- Canny algorithm -------------------------------------------------------------
  
  canny :: Float -> Float -> Acc (Image RGBA32) -> (Acc (Image Float), Acc (Vector Int))
  canny (constant -> low) (constant -> high)
    = stage1
    . nonMaximumSuppression low high
    . gradientMagDir low
    . gaussianY
    . gaussianX
    . toGreyscale
    where
      stage1 x = (x, selectStrong x)
  
  
  -- Accelerate component --------------------------------------------------------
  
  type Image a            = Array DIM2 a
  
  type Stencil5x1 a       = (Stencil3 a, Stencil5 a, Stencil3 a)
  type Stencil1x5 a       = (Stencil3 a, Stencil3 a, Stencil3 a, Stencil3 a, Stencil3 a)
  
  -- Classification of the output pixel
  data Orient     = Undef | PosD | Vert | NegD | Horiz
  data Edge       = None  | Weak | Strong
  
  orient :: Orient -> Int
  orient Undef    = 0
  orient PosD     = 64
  orient Vert     = 128
  orient NegD     = 192
  orient Horiz    = 255
  
  orient' :: Orient -> Exp Int
  orient' = constant . orient
  
  edge :: Edge -> Float
  edge None       = 0
  edge Weak       = 0.5
  edge Strong     = 1.0
  
  edge' :: Edge -> Exp Float
  edge' = constant . edge
  
  convolve5x1 :: Num a => [Exp a] -> Stencil5x1 a -> Exp a
  convolve5x1 kernel (_, (a,b,c,d,e), _)
    = P.sum $ P.zipWith (*) kernel [a,b,c,d,e]
  
  convolve1x5 :: Num a => [Exp a] -> Stencil1x5 a -> Exp a
  convolve1x5 kernel ((_,a,_), (_,b,_), (_,c,_), (_,d,_), (_,e,_))
    = P.sum $ P.zipWith (*) kernel [a,b,c,d,e]
  
  
  -- RGB to Greyscale conversion, in the range [0,255]
  --
  toGreyscale :: Acc (Image RGBA32) -> Acc (Image Float)
  toGreyscale = map (\rgba -> 255 * luminance (unpackRGB rgba))
  
  
  -- Separable Gaussian blur in the x- and y-directions
  --
  gaussianX :: Acc (Image Float) -> Acc (Image Float)
  gaussianX = stencil (convolve5x1 gaussian) A.clamp
    where
      gaussian = P.map (/16) [ 1, 4, 6, 4, 1 ]
  
  gaussianY :: Acc (Image Float) -> Acc (Image Float)
  gaussianY = stencil (convolve1x5 gaussian) A.clamp
    where
      gaussian = P.map (/16) [ 1, 4, 6, 4, 1 ]
  
  
  -- Gradients in the x- and y- directions
  --
  gradientX :: Acc (Image Float) -> Acc (Image Float)
  gradientX = stencil grad A.clamp
    where
      grad :: Stencil3x3 Float -> Exp Float
      grad ((u, _, x)
           ,(v, _, y)
           ,(w, _, z)) = x + (2*y) + z - u - (2*v) - w
  
  gradientY :: Acc (Image Float) -> Acc (Image Float)
  gradientY = stencil grad A.clamp
    where
      grad :: Stencil3x3 Float -> Exp Float
      grad ((x, y, z)
           ,(_, _, _)
           ,(u, v, w)) = x + (2*y) + z - u - (2*v) - w
  
  
  -- Classify the magnitude and orientation of the image gradient.
  --
  -- Because accelerate supports generalised stencil functions, not just
  -- convolutions, we can combine the x- and y- sobel operators and save some
  -- memory bandwidth.
  --
  gradientMagDir
      :: Exp Float
      -> Acc (Image Float)
      -> Acc (Array DIM2 (Float,Int))
  gradientMagDir low = stencil magdir A.clamp
    where
      magdir :: Stencil3x3 Float -> Exp (Float,Int)
      magdir ((v0, v1, v2)
             ,(v3,  _, v4)
             ,(v5, v6, v7)) =
        let
            -- Image gradients
            dx          = v2 + (2*v4) + v7 - v0 - (2*v3) - v5
            dy          = v0 + (2*v1) + v2 - v5 - (2*v6) - v7
  
            -- Magnitude
            mag         = sqrt (dx * dx + dy * dy)
  
            -- Direction
            --
            -- Determine the angle of the vector and rotate it around a bit to
            -- make the segments easier to classify
            theta       = atan2 dy dx
            alpha       = (theta - (pi/8)) * (4/pi)
  
            -- Normalise the angle to between [0..8)
            norm        = alpha + 8 * fromIntegral (boolToInt (alpha <= 0))
  
            -- Try to avoid doing explicit tests, to avoid warp divergence
            undef       = abs dx <= low && abs dy <= low
            dir         = boolToInt (not undef) * ((64 * (1 + floor norm `mod` 4)) `min` 255)
        in
        lift (mag, dir)
  
  
  -- Non-maximum suppression classifies pixels that are the local maximum along
  -- the direction of the image gradient as either strong or weak edges. All other
  -- pixels are not considered edges at all.
  --
  -- The image intensity is in the range [0,1]
  --
  nonMaximumSuppression
    :: Exp Float
    -> Exp Float
    -> Acc (Image (Float,Int))
    -> Acc (Image Float)
  nonMaximumSuppression low high magdir =
    generate (shape magdir) $ \ix ->
      let -- The input parameters
          --
          (mag, dir)      = unlift (magdir ! ix)
          Z :. h :. w     = unlift (shape magdir)
          Z :. y :. x     = unlift ix
  
          -- Determine the points that lie either side of this point along to the
          -- direction of the image gradient.
          --
          -- The direction coding:
          --
          --   192   128   64
          --          |
          --   255 --- ---
          --
          offsetx         = dir > orient' Vert  ? (-1, dir < orient' Vert ? (1, 0))
          offsety         = dir < orient' Horiz ? (-1, 0)
  
          (fwd, _)        = unlift $ magdir ! lift (limit (Z :. y+offsety :. x+offsetx)) :: (Exp Float, Exp Int)
          (rev, _)        = unlift $ magdir ! lift (limit (Z :. y-offsety :. x-offsetx)) :: (Exp Float, Exp Int)
  
          limit (Z:.u:.v) = Z :. 0 `max` u `min` (h-1) :. 0 `max` v `min` (w-1)
  
          -- Try to avoid doing explicit tests to avoid warp divergence.
          --
          none            = dir == orient' Undef || mag < low || mag < fwd || mag < rev
          strong          = mag >= high
      in
      fromIntegral (boolToInt (not none) * (1 + boolToInt strong)) * 0.5
  
  
  -- Extract the linear indices of the strong edges
  --
  selectStrong
    :: Acc (Image Float)
    -> Acc (Array DIM1 Int)
  selectStrong img =
    let strong            = map (\x -> boolToInt (x == edge' Strong)) (flatten img)
        (targetIdx, len)  = unlift (scanl' (+) 0 strong)
        indices           = enumFromN (index1 $ size img) 0
        zeros             = fill (index1 $ the len) 0
    in
    permute const zeros (\ix -> strong!ix == 0 ? (Nothing_, Just_ (I1 (targetIdx!ix)))) indices
  
  