{-# LANGUAGE ForeignFunctionInterface, TemplateHaskell, QuasiQuotes, BangPatterns #-} -- | Does Canny Edge Detection Processing. -- Needs Repa 2.0.0.1 -- For best performance this module needs to be compiled like: -- -- ghc -threaded -rtsopts -Odph -fllvm -optlo-O3 -- -fno-liberate-case -funfolding-use-threshold100 -funfolding-keeness-factor100 -- -c Process.hs -- module Process where import Foreign.Ptr import Foreign.C.Types import Foreign.Storable import Data.Word import Data.Array.Repa as R import Data.Array.Repa.ByteString as R import Data.Array.Repa.Stencil as R import Data.Array.Repa.Specialised.Dim2 as R import Control.Monad import System.IO.Unsafe import qualified Data.Vector.Unboxed.Mutable as VM import qualified Data.Vector.Unboxed as V foreign export ccall processImage :: Int -> Int -> Int -> Float -> Float -> Int -> Int -> Ptr Word8 -> IO () processImage :: Int -- output select -> Int -- enable blur -> Int -- enable invert -> Float -- threshold low -> Float -- threshold high -> Int -- image width -> Int -- image height -> Ptr Word8 -- ptr to image data -> IO () processImage output enableBlur enableInvert threshLow threshHigh width height ptr = do -- copy data out of the buffer arr <- liftM force $ R.copyFromPtrWord8 (Z :. height :. width :. (4 :: Int)) ptr let config = Config { configEnableBlur = enableBlur == 1 , configEnableInvert = enableInvert == 1 , configThresholdLow = threshLow , configThresholdHigh = threshHigh , configOutputSelect = [ OutputGrey, OutputBlur, OutputDX, OutputDY, OutputMag , OutputOrient, OutputMax, OutputLink] !! output } -- transform the image let arr' = force $ fromGreyScale (configEnableInvert config) $ edgeDetect config $ toGreyScale arr -- copy data back over the original buffer R.copyToPtrWord8 ptr arr' -- Edge Detect -------------------------------------------------------------------------- -- | We mostly work with Float images. type Image a = Array DIM2 a -- | What stage of the algorithm we want to return on the output. data OutputSelect = OutputGrey -- 0 | OutputBlur -- 1 | OutputDX -- 2 | OutputDY -- 3 | OutputMag -- 4 | OutputOrient -- 5 | OutputMax -- 6 | OutputLink -- 7 -- | Configuration for edge detector. data Config = Config { configEnableBlur :: Bool , configEnableInvert :: Bool , configThresholdLow :: Float , configThresholdHigh :: Float , configOutputSelect :: OutputSelect } edgeDetect :: Config -> Image Float -> Image Float {-# NOINLINE edgeDetect #-} edgeDetect config arrInput = let -- Blur input image. arrBluredX = blurSepX arrInput arrBlured = blurSepY arrBluredX arrBluredOut = if configEnableBlur config then arrBlured else arrInput -- Compute magnitude and orientation of vector gradient arrDX = gradientX arrBluredOut arrDY = gradientY arrBluredOut arrMagOrient = gradientMagOrient (configThresholdLow config) arrDX arrDY -- Non-maximum suppression, mark points as weak, strong, or no edge. arrSuppress = suppress (configThresholdLow config) (configThresholdHigh config) arrMagOrient -- Collect up points marked as strong edges. arrStrong = selectStrong arrSuppress -- Mark strong edges, and weak ones connected to strong ones. arrLink = wildfire arrSuppress arrStrong in case configOutputSelect config of OutputGrey -> arrInput OutputBlur -> arrBlured OutputDX -> force2 $ R.map abs arrDX OutputDY -> force2 $ R.map abs arrDY OutputMag -> force2 $ R.map fst arrMagOrient OutputOrient -> force2 $ R.map snd arrMagOrient OutputMax -> arrSuppress OutputLink -> arrLink -- Constants ------------------------------------------------------------------ -- Classification of the direction of the image gradient. orientUndef = 0 :: Float orientPosDiag = 64 :: Float orientVert = 128 :: Float orientNegDiag = 192 :: Float orientHoriz = 255 :: Float -- Classification of the output pixel. data Edge = None | Weak | Strong edge None = 0 :: Float edge Weak = 128 :: Float edge Strong = 255 :: Float ------------------------------------------------------------------------------- -- | RGB to greyscale conversion. {-# NOINLINE toGreyScale #-} toGreyScale :: Array DIM3 Word8 -> Image Float toGreyScale arr@(Array _ [Region RangeAll (GenManifest _)]) = arr `deepSeqArray` force2 $ traverse arr (\(sh :. _) -> sh) (\get ix -> rgbToLuminance (get (ix :. 0)) (get (ix :. 1)) (get (ix :. 2))) where {-# INLINE rgbToLuminance #-} rgbToLuminance :: Word8 -> Word8 -> Word8 -> Float rgbToLuminance r g b = floatOfWord8 r * 0.3 + floatOfWord8 g * 0.59 + floatOfWord8 b * 0.11 {-# INLINE floatOfWord8 #-} floatOfWord8 :: Word8 -> Float floatOfWord8 w8 = fromIntegral (fromIntegral w8 :: Int) -- | Separable Gaussian blur in the X direction. {-# NOINLINE blurSepX #-} blurSepX :: Image Float -> Image Float blurSepX arr@(Array _ [Region RangeAll (GenManifest _)]) = arr `deepSeqArray` force2 $ forStencil2 BoundClamp arr [stencil2| 1 4 6 4 1 |] -- | Separable Gaussian blur in the Y direction. {-# NOINLINE blurSepY #-} blurSepY :: Image Float -> Image Float blurSepY arr@(Array _ [Region RangeAll (GenManifest _)]) = arr `deepSeqArray` force2 $ R.map (/ 256) $ forStencil2 BoundClamp arr [stencil2| 1 4 6 4 1 |] -- | Compute gradient in the X direction. {-# NOINLINE gradientX #-} gradientX :: Image Float -> Image Float gradientX img@(Array _ [Region RangeAll (GenManifest _)]) = img `deepSeqArray` force2 $ forStencil2 (BoundConst 0) img [stencil2| -1 0 1 -2 0 2 -1 0 1 |] -- | Compute gradient in the Y direction. {-# NOINLINE gradientY #-} gradientY :: Image Float -> Image Float gradientY img@(Array _ [Region RangeAll (GenManifest _)]) = img `deepSeqArray` force2 $ forStencil2 (BoundConst 0) img [stencil2| 1 2 1 0 0 0 -1 -2 -1 |] -- | Classify the magnitude and orientation of the vector gradient. {-# NOINLINE gradientMagOrient #-} gradientMagOrient :: Float -> Image Float -> Image Float -> Image (Float, Float) gradientMagOrient !threshLow dX@(Array _ [Region RangeAll (GenManifest _)]) dY@(Array _ [Region RangeAll (GenManifest _)]) = [dX, dY] `deepSeqArrays` force2 $ R.zipWith magOrient dX dY where {-# INLINE magOrient #-} magOrient :: Float -> Float -> (Float, Float) magOrient !x !y = (magnitude x y, orientation x y) {-# INLINE magnitude #-} magnitude :: Float -> Float -> Float magnitude !x !y = sqrt (x * x + y * y) {-# INLINE orientation #-} orientation :: Float -> Float -> Float orientation !x !y -- Don't bother computing orientation if vector is below threshold. -- We wont' use it in the next step anyway. | 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. -- Though this could probably be improved scaling to an integer and using `mod`. in if dNorm >= 4 then if dNorm >= 6 then if dNorm >= 7 then orientHoriz -- 7 else orientNegDiag -- 6 else if dNorm >= 5 then orientVert -- 5 else orientPosDiag -- 4 else if dNorm >= 2 then if dNorm >= 3 then orientHoriz -- 3 else orientNegDiag -- 2 else if dNorm >= 1 then orientVert -- 1 else orientPosDiag -- 0 -- | Suppress pixels which are not local maxima, and use the magnitude to classify maxima -- into strong and weak (potential) edges. {-# NOINLINE suppress #-} suppress :: Float -> Float -> Image (Float, Float) -> Image Float suppress threshLow threshHigh dMagOrient@(Array shSrc [Region RangeAll (GenManifest _)]) = dMagOrient `deepSeqArray` force2 $ makeBordered2 shSrc 1 (GenCursor id addDim (const 0)) (GenCursor id addDim compare) where {-# INLINE compare #-} compare 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 -- | Select indices of strong edges. {-# NOINLINE selectStrong #-} selectStrong :: Image Float -> Array DIM1 Int selectStrong img@(Array _ [Region RangeAll (GenManifest vec)]) = img `deepSeqArray` let {-# INLINE match #-} match ix = vec `V.unsafeIndex` ix == edge Strong {-# INLINE process #-} process ix = ix in select match process (size $ extent img) -- | Trace out strong edges in the final image. -- Also trace out weak edges that are connected to strong edges. {-# NOINLINE wildfire #-} wildfire :: Image Float -- ^ Image with strong and weak edges set. -> Array DIM1 Int -- ^ Array containing flat indices of strong edges. -> Image Float wildfire img@(Array _ [Region RangeAll (GenManifest _)]) arrStrong = unsafePerformIO $ do (sh, vec) <- wildfireIO return $ sh `seq` vec `seq` Array sh [Region RangeAll (GenManifest vec)] where lenImg = R.size $ R.extent img lenStrong = R.size $ R.extent arrStrong vStrong = toVector 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 Float -> 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 -- | Convert back from greyscale image fromGreyScale :: Bool -> Image Float -> Array DIM3 Word8 {-# NOINLINE fromGreyScale #-} fromGreyScale False arr@(Array _ [Region RangeAll (GenManifest _)]) = arr `deepSeqArray` force $ unsafeTraverse arr (\sh -> sh :. (4 :: Int)) (\get (ix :. comp) -> case comp of 0 -> truncate (get ix) 1 -> truncate (get ix) 2 -> truncate (get ix) 3 -> truncate (get ix)) fromGreyScale True arr@(Array _ [Region RangeAll (GenManifest _)]) = arr `deepSeqArray` force $ unsafeTraverse arr (\sh -> sh :. (4 :: Int)) (\get (ix :. comp) -> case comp of 0 -> 255 - truncate (get ix) 1 -> 255 - truncate (get ix) 2 -> 255 - truncate (get ix) 3 -> 255 - truncate (get ix))