{-# 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 [ ]" , " " ] 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 #-}