repa-examples-3.4.1.1: examples/Canny/src-repa/Main.hs
{-# LANGUAGE PackageImports, BangPatterns, QuasiQuotes, PatternGuards,
MagicHash, ScopedTypeVariables, TypeFamilies #-}
{-# OPTIONS -Wall -fno-warn-missing-signatures -fno-warn-incomplete-patterns #-}
-- | Canny edge detector.
--
-- NOTE: for best performance this needs to be compiled with the following GHC options:
-- -fllvm -optlo-O3 -Odph -fno-liberate-case
-- -funfolding-use-threshold100 -funfolding-keeness-factor100
--
import Data.List
import Data.Word
import Data.Int
import Control.Monad
import System.Environment
import Data.Array.Repa as R
import Data.Array.Repa.Repr.Unboxed as U
import Data.Array.Repa.Repr.Cursored as C
import Data.Array.Repa.Stencil
import Data.Array.Repa.Stencil.Dim2
import Data.Array.Repa.Specialised.Dim2
import Data.Array.Repa.Algorithms.Pixel
import Data.Array.Repa.IO.BMP
import Data.Array.Repa.IO.Timing
import Debug.Trace
import GHC.Exts
import qualified Data.Vector.Unboxed.Mutable as VM
import qualified Data.Vector.Unboxed as V
import qualified Prelude as P
import Prelude hiding (compare)
type Image a = Array U DIM2 a
-- Constants ------------------------------------------------------------------
orientUndef = 0 :: Word8
orientPosDiag = 64 :: Word8
orientVert = 128 :: Word8
orientNegDiag = 192 :: Word8
orientHoriz = 255 :: Word8
data Edge = None | Weak | Strong
edge None = 0 :: Word8
edge Weak = 128 :: Word8
edge Strong = 255 :: Word8
-- Main routine ---------------------------------------------------------------
main
= do args <- getArgs
case args of
[fileIn, fileOut]
-> run 0 50 100 fileIn fileOut
[loops, threshLow, threshHigh, fileIn, fileOut]
-> run (read loops) (read threshLow) (read threshHigh) fileIn fileOut
_ -> putStrLn
$ concat [ "repa-canny [<loops::Int> <threshLow::Int> <threshHigh::Int>]"
, " <fileIn.bmp> <fileOut.bmp>" ]
run loops threshLow threshHigh fileIn fileOut
= do arrInput <- liftM (either (error . show) id)
$ readImageFromBMP fileIn
(arrResult, tTotal)
<- time $ process loops threshLow threshHigh arrInput
when (loops >= 1)
$ putStrLn $ "\nTOTAL\n"
putStr $ prettyTime tTotal
writeImageToBMP fileOut (U.zip3 arrResult arrResult arrResult)
process loops threshLow threshHigh arrInput
= do arrGrey <- timeStage loops "toGreyScale"
$ toGreyScale arrInput
arrBluredX <- timeStage loops "blurX"
$ blurSepX arrGrey
arrBlured <- timeStage loops "blurY"
$ blurSepY arrBluredX
arrDX <- timeStage loops "diffX"
$ gradientX arrBlured
arrDY <- timeStage loops "diffY"
$ gradientY arrBlured
arrMagOrient <- timeStage loops "magOrient"
$ gradientMagOrient threshLow arrDX arrDY
arrSuppress <- timeStage loops "suppress"
$ suppress threshLow threshHigh arrMagOrient
arrStrong <- timeStage loops "select"
$ selectStrong arrSuppress
arrEdges <- timeStage loops "wildfire"
$ wildfire arrSuppress arrStrong
return arrEdges
-- | Wrapper to time each stage of the algorithm.
timeStage
:: Int
-> String
-> IO (Array U sh a)
-> IO (Array U sh a)
timeStage loops name fn
= do
let burn !n
= do !arr <- fn
if n <= 1 then return arr
else burn (n - 1)
traceEventIO $ "**** Stage " P.++ name P.++ " begin."
(arrResult, t)
<- time $ do !arrResult' <- burn loops
return arrResult'
traceEventIO $ "**** Stage " P.++ name P.++ " end."
when (loops >= 1)
$ putStr $ name P.++ "\n"
P.++ unlines [ " " P.++ l | l <- lines $ prettyTime t ]
return arrResult
{-# NOINLINE timeStage #-}
-------------------------------------------------------------------------------
-- | RGB to greyscale conversion.
toGreyScale :: Image (Word8, Word8, Word8) -> IO (Image Float)
toGreyScale arr
= computeP
$ R.map (* 255)
$ R.map floatLuminanceOfRGB8 arr
{-# NOINLINE toGreyScale #-}
-- | Separable Gaussian blur in the X direction.
blurSepX :: Image Float -> IO (Image Float)
blurSepX arr
= computeP
$ forStencil2 BoundClamp arr
[stencil2| 1 4 6 4 1 |]
{-# NOINLINE blurSepX #-}
-- | Separable Gaussian blur in the Y direction.
blurSepY :: Image Float -> IO (Image Float)
blurSepY arr
= computeP
$ R.smap (/ 256)
$ forStencil2 BoundClamp arr
[stencil2| 1
4
6
4
1 |]
{-# NOINLINE blurSepY #-}
-- | Compute gradient in the X direction.
gradientX :: Image Float -> IO (Image Float)
gradientX img
= computeP
$ forStencil2 BoundClamp img
[stencil2| -1 0 1
-2 0 2
-1 0 1 |]
{-# NOINLINE gradientX #-}
-- | Compute gradient in the Y direction.
gradientY :: Image Float -> IO (Image Float)
gradientY img
= computeP
$ forStencil2 BoundClamp img
[stencil2| 1 2 1
0 0 0
-1 -2 -1 |]
{-# NOINLINE gradientY #-}
-- | Classify the magnitude and orientation of the vector gradient.
gradientMagOrient
:: Float -> Image Float -> Image Float -> IO (Image (Float, Word8))
gradientMagOrient !threshLow dX dY
= computeP
$ R.zipWith magOrient dX dY
where magOrient :: Float -> Float -> (Float, Word8)
magOrient !x !y
= (magnitude x y, orientation x y)
{-# INLINE magOrient #-}
magnitude :: Float -> Float -> Float
magnitude !x !y
= sqrt (x * x + y * y)
{-# INLINE magnitude #-}
{-# INLINE orientation #-}
orientation :: Float -> Float -> Word8
orientation !x !y
-- Don't bother computing orientation if vector is below threshold.
| x >= negate threshLow, x < threshLow
, y >= negate threshLow, y < threshLow
= orientUndef
| otherwise
= let -- Determine the angle of the vector and rotate it around a bit
-- to make the segments easier to classify.
!d = atan2 y x
!dRot = (d - (pi/8)) * (4/pi)
-- Normalise angle to beween 0..8
!dNorm = if dRot < 0 then dRot + 8 else dRot
-- Doing explicit tests seems to be faster than using the FP floor function.
in fromIntegral
$ I# (if dNorm >= 4
then if dNorm >= 6
then if dNorm >= 7
then 255# -- 7
else 192# -- 6
else if dNorm >= 5
then 128# -- 5
else 64# -- 4
else if dNorm >= 2
then if dNorm >= 3
then 255# -- 3
else 192# -- 2
else if dNorm >= 1
then 128# -- 1
else 64#) -- 0
{-# NOINLINE gradientMagOrient #-}
-- | Suppress pixels that are not local maxima, and use the magnitude to classify maxima
-- into strong and weak (potential) edges.
suppress :: Float -> Float -> Image (Float, Word8) -> IO (Image Word8)
suppress !threshLow !threshHigh !dMagOrient
= computeP
$ makeBordered2
(extent dMagOrient) 1
(makeCursored (extent dMagOrient) id addDim comparePts)
(fromFunction (extent dMagOrient) (const 0))
where {-# INLINE comparePts #-}
comparePts d@(sh :. i :. j)
| o == orientUndef = edge None
| o == orientHoriz = isMax (getMag (sh :. i :. j-1)) (getMag (sh :. i :. j+1))
| o == orientVert = isMax (getMag (sh :. i-1 :. j)) (getMag (sh :. i+1 :. j))
| o == orientNegDiag = isMax (getMag (sh :. i-1 :. j-1)) (getMag (sh :. i+1 :. j+1))
| o == orientPosDiag = isMax (getMag (sh :. i-1 :. j+1)) (getMag (sh :. i+1 :. j-1))
| otherwise = edge None
where
!o = getOrient d
!m = getMag (Z :. i :. j)
getMag = fst . (R.unsafeIndex dMagOrient)
getOrient = snd . (R.unsafeIndex dMagOrient)
{-# INLINE isMax #-}
isMax !intensity1 !intensity2
| m < threshLow = edge None
| m < intensity1 = edge None
| m < intensity2 = edge None
| m < threshHigh = edge Weak
| otherwise = edge Strong
{-# NOINLINE suppress #-}
-- | Select indices of strong edges.
-- TODO: If would better if we could medge this into the above stage, and
-- record the strong edge during non-maximum suppression, but Repa
-- doesn't provide a fused mapFilter primitive yet.
selectStrong :: Image Word8 -> IO (Array U DIM1 Int)
selectStrong img
= let vec = toUnboxed img
match ix = vec `V.unsafeIndex` ix == edge Strong
{-# INLINE match #-}
process' ix = ix
{-# INLINE process' #-}
in selectP match process' (size $ extent img)
{-# NOINLINE selectStrong #-}
-- | Trace out strong edges in the final image.
-- Also trace out weak edges that are connected to strong edges.
wildfire
:: Image Word8 -- ^ Image with strong and weak edges set.
-> Array U DIM1 Int -- ^ Array containing flat indices of strong edges.
-> IO (Image Word8)
wildfire img arrStrong
= do (sh, vec) <- wildfireIO
return $ sh `seq` vec `seq` fromUnboxed sh vec
where lenImg = R.size $ R.extent img
lenStrong = R.size $ R.extent arrStrong
vStrong = toUnboxed arrStrong
wildfireIO
= do -- Stack of image indices we still need to consider.
vStrong' <- V.thaw vStrong
vStack <- VM.grow vStrong' (lenImg - lenStrong)
-- Burn in new edges.
vImg <- VM.unsafeNew lenImg
VM.set vImg 0
burn vImg vStack lenStrong
vImg' <- V.unsafeFreeze vImg
return (extent img, vImg')
burn :: VM.IOVector Word8 -> VM.IOVector Int -> Int -> IO ()
burn !vImg !vStack !top
| top == 0
= return ()
| otherwise
= do let !top' = top - 1
n <- VM.unsafeRead vStack top'
let (Z :. y :. x) = fromIndex (R.extent img) n
let {-# INLINE push #-}
push t = pushWeak vImg vStack t
VM.write vImg n (edge Strong)
>> push (Z :. y - 1 :. x - 1) top'
>>= push (Z :. y - 1 :. x )
>>= push (Z :. y - 1 :. x + 1)
>>= push (Z :. y :. x - 1)
>>= push (Z :. y :. x + 1)
>>= push (Z :. y + 1 :. x - 1)
>>= push (Z :. y + 1 :. x )
>>= push (Z :. y + 1 :. x + 1)
>>= burn vImg vStack
-- If this ix is weak in the source then set it to strong in the
-- result and push the ix onto the stack.
{-# INLINE pushWeak #-}
pushWeak vImg vStack ix top
= do let n = toIndex (extent img) ix
xDst <- VM.unsafeRead vImg n
let xSrc = img `R.unsafeIndex` ix
if xDst == edge None
&& xSrc == edge Weak
then do
VM.unsafeWrite vStack top (toIndex (extent img) ix)
return (top + 1)
else return top
{-# NOINLINE wildfire #-}