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

-- | Provides high level functions to do geometric transformations on images.
--
-- Every transformation is been declared @INLINABLE@ so new image types could be
-- specialized.
module Vision.Image.Transform (
      InterpolMethod (..), crop, resize, horizontalFlip, verticalFlip, floodFill
    ) where

import Control.Monad (when)
import Control.Monad.Primitive (PrimMonad (..))
import Data.RatioInt (RatioInt, (%))

import Vision.Image.Class (
      MaskedImage (..), Image (..), ImageChannel, FromFunction (..), (!)
    )
import Vision.Image.Interpolate (Interpolable, bilinearInterpol)
import Vision.Image.Mutable (MutableImage (..))
import Vision.Primitive (
      Z (..), (:.) (..), Point, RPoint (..), Rect (..), Size, ix2, toLinearIndex
    )

-- | Defines the set of possible methods for pixel interpolations when looking
-- for a pixel at floating point coordinates.
data InterpolMethod =
      TruncateInteger -- ^ Selects the top left pixel (fastest).
    | NearestNeighbor -- ^ Selects the nearest pixel (fast).
    | Bilinear        -- ^ Does a double linear interpolation over the four
                      -- surrounding points (slow).

-- | Maps the content of the image\'s rectangle in a new image.
crop :: (Image i1, FromFunction i2, ImagePixel i1 ~ FromFunctionPixel i2)
     => Rect -> i1 -> i2
crop :: forall i1 i2.
(Image i1, FromFunction i2,
 ImagePixel i1 ~ FromFunctionPixel i2) =>
Rect -> i1 -> i2
crop !(Rect Int
rx Int
ry Int
rw Int
rh) !i1
img =
    forall i.
FromFunction i =>
Size -> (Size -> FromFunctionPixel i) -> i
fromFunction (DIM0
Z forall tail head. tail -> head -> tail :. head
:. Int
rh forall tail head. tail -> head -> tail :. head
:. Int
rw) forall a b. (a -> b) -> a -> b
$ \(DIM0
Z :. Int
y :. Int
x) ->
        i1
img forall i. Image i => i -> Size -> ImagePixel i
! Int -> Int -> Size
ix2 (Int
ry forall a. Num a => a -> a -> a
+ Int
y) (Int
rx forall a. Num a => a -> a -> a
+ Int
x)
{-# INLINABLE crop #-}

-- | Resizes the 'Image' using the given interpolation method.
resize :: (Image i1, Interpolable (ImagePixel i1), FromFunction i2
         , ImagePixel i1 ~ FromFunctionPixel i2, Integral (ImageChannel i1))
       => InterpolMethod -> Size -> i1 -> i2
resize :: forall i1 i2.
(Image i1, Interpolable (ImagePixel i1), FromFunction i2,
 ImagePixel i1 ~ FromFunctionPixel i2,
 Integral (ImageChannel i1)) =>
InterpolMethod -> Size -> i1 -> i2
resize !InterpolMethod
method !size' :: Size
size'@(DIM0
Z :. Int
h' :. Int
w') !i1
img =
    case InterpolMethod
method of
        InterpolMethod
TruncateInteger ->
            let !widthRatio :: Double
widthRatio   = forall a. Integral a => a -> Double
double Int
w forall a. Fractional a => a -> a -> a
/ forall a. Integral a => a -> Double
double Int
w'
                !heightRatio :: Double
heightRatio  = forall a. Integral a => a -> Double
double Int
h forall a. Fractional a => a -> a -> a
/ forall a. Integral a => a -> Double
double Int
h'
                line :: Int -> Int
line !Int
y' = forall a b. (RealFrac a, Integral b) => a -> b
truncate forall a b. (a -> b) -> a -> b
$ (forall a. Integral a => a -> Double
double Int
y' forall a. Num a => a -> a -> a
+ Double
0.5) forall a. Num a => a -> a -> a
* Double
heightRatio forall a. Num a => a -> a -> a
- Double
0.5
                {-# INLINE line #-}
                col :: Int -> Int
col  !Int
x' = forall a b. (RealFrac a, Integral b) => a -> b
truncate forall a b. (a -> b) -> a -> b
$ (forall a. Integral a => a -> Double
double Int
x' forall a. Num a => a -> a -> a
+ Double
0.5) forall a. Num a => a -> a -> a
* Double
widthRatio  forall a. Num a => a -> a -> a
- Double
0.5
                {-# INLINE col #-}
                f :: Int -> Size -> ImagePixel i1
f !Int
y !(DIM0
Z :. Int
_ :. Int
x') = let !x :: Int
x = Int -> Int
col Int
x'
                                       in i1
img forall i. Image i => i -> Size -> ImagePixel i
! Int -> Int -> Size
ix2 Int
y Int
x
                {-# INLINE f #-}
            in forall i a.
FromFunction i =>
Size -> (Int -> a) -> (a -> Size -> FromFunctionPixel i) -> i
fromFunctionLine Size
size' Int -> Int
line Int -> Size -> ImagePixel i1
f
        InterpolMethod
NearestNeighbor ->
            let !widthRatio :: Double
widthRatio   = forall a. Integral a => a -> Double
double Int
w forall a. Fractional a => a -> a -> a
/ forall a. Integral a => a -> Double
double Int
w'
                !heightRatio :: Double
heightRatio  = forall a. Integral a => a -> Double
double Int
h forall a. Fractional a => a -> a -> a
/ forall a. Integral a => a -> Double
double Int
h'
                line :: Int -> Int
line !Int
y' = forall a b. (RealFrac a, Integral b) => a -> b
round forall a b. (a -> b) -> a -> b
$ (forall a. Integral a => a -> Double
double Int
y' forall a. Num a => a -> a -> a
+ Double
0.5) forall a. Num a => a -> a -> a
* Double
heightRatio forall a. Num a => a -> a -> a
- Double
0.5
                {-# INLINE line #-}
                col :: Int -> Int
col  !Int
x' = forall a b. (RealFrac a, Integral b) => a -> b
round forall a b. (a -> b) -> a -> b
$ (forall a. Integral a => a -> Double
double Int
x' forall a. Num a => a -> a -> a
+ Double
0.5) forall a. Num a => a -> a -> a
* Double
widthRatio  forall a. Num a => a -> a -> a
- Double
0.5
                {-# INLINE col #-}
                f :: Int -> Size -> ImagePixel i1
f !Int
y !(DIM0
Z :. Int
_ :. Int
x') = let !x :: Int
x = Int -> Int
col Int
x'
                                       in i1
img forall i. Image i => i -> Size -> ImagePixel i
! Int -> Int -> Size
ix2 Int
y Int
x
                {-# INLINE f #-}
            in forall i a.
FromFunction i =>
Size -> (Int -> a) -> (a -> Size -> FromFunctionPixel i) -> i
fromFunctionLine Size
size' Int -> Int
line Int -> Size -> ImagePixel i1
f
        InterpolMethod
Bilinear ->
            let !widthRatio :: RatioInt
widthRatio  = Int
w Int -> Int -> RatioInt
% Int
w'
                !maxWidth :: RatioInt
maxWidth = forall a. Integral a => a -> RatioInt
ratio (Int
w forall a. Num a => a -> a -> a
- Int
1)
                !heightRatio :: RatioInt
heightRatio = (Int
h forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> RatioInt
% (Int
h' forall a. Num a => a -> a -> a
- Int
1)
                !maxHeight :: RatioInt
maxHeight = forall a. Integral a => a -> RatioInt
ratio (Int
h forall a. Num a => a -> a -> a
- Int
1)

                -- Limits the interpolation to inner pixel as first and last
                -- pixels can have out of bound coordinates.
                bound :: c -> c -> c
bound !c
limit = forall a. Ord a => a -> a -> a
min c
limit forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => a -> a -> a
max c
0
                {-# INLINE bound #-}

                line :: Int -> RatioInt
line !Int
y' = forall {c}. (Ord c, Num c) => c -> c -> c
bound RatioInt
maxHeight forall a b. (a -> b) -> a -> b
$   (forall a. Integral a => a -> RatioInt
ratio Int
y' forall a. Num a => a -> a -> a
+ RatioInt
0.5) forall a. Num a => a -> a -> a
* RatioInt
heightRatio
                                             forall a. Num a => a -> a -> a
- RatioInt
0.5
                {-# INLINE line #-}
                col :: Int -> RatioInt
col  !Int
x' = forall {c}. (Ord c, Num c) => c -> c -> c
bound RatioInt
maxWidth  forall a b. (a -> b) -> a -> b
$   (forall a. Integral a => a -> RatioInt
ratio Int
x' forall a. Num a => a -> a -> a
+ RatioInt
0.5) forall a. Num a => a -> a -> a
* RatioInt
widthRatio
                                             forall a. Num a => a -> a -> a
- RatioInt
0.5
                {-# INLINE col #-}
                f :: RatioInt -> RatioInt -> Size -> ImagePixel i1
f !RatioInt
y !RatioInt
x Size
_ = i1
img forall i.
(Image i, Interpolable (ImagePixel i),
 Integral (ImageChannel i)) =>
i -> RPoint -> ImagePixel i
`bilinearInterpol` RatioInt -> RatioInt -> RPoint
RPoint RatioInt
x RatioInt
y
                {-# INLINE f #-}
            in forall i b a.
(FromFunction i, Storable b) =>
Size
-> (Int -> a)
-> (Int -> b)
-> (a -> b -> Size -> FromFunctionPixel i)
-> i
fromFunctionCached Size
size' Int -> RatioInt
line Int -> RatioInt
col RatioInt -> RatioInt -> Size -> ImagePixel i1
f
  where
    !(DIM0
Z :. Int
h :. Int
w) = forall i. MaskedImage i => i -> Size
shape i1
img
{-# INLINABLE resize #-}

-- | Reverses the image horizontally.
horizontalFlip :: (Image i1, FromFunction i2
                  , ImagePixel i1 ~ FromFunctionPixel i2)
               => i1 -> i2
horizontalFlip :: forall i1 i2.
(Image i1, FromFunction i2,
 ImagePixel i1 ~ FromFunctionPixel i2) =>
i1 -> i2
horizontalFlip !i1
img =
    let f :: Size -> ImagePixel i1
f !(DIM0
Z :. Int
y :. Int
x') = let !x :: Int
x = Int
maxX forall a. Num a => a -> a -> a
- Int
x'
                            in i1
img forall i. Image i => i -> Size -> ImagePixel i
! Int -> Int -> Size
ix2 Int
y Int
x
        {-# INLINE f #-}
    in forall i.
FromFunction i =>
Size -> (Size -> FromFunctionPixel i) -> i
fromFunction Size
size Size -> ImagePixel i1
f
  where
    !size :: Size
size@(DIM0
Z :. Int
_ :. Int
w) = forall i. MaskedImage i => i -> Size
shape i1
img
    !maxX :: Int
maxX = Int
w forall a. Num a => a -> a -> a
- Int
1
{-# INLINABLE horizontalFlip #-}

-- | Reverses the image vertically.
verticalFlip :: (Image i1, FromFunction i2
                , ImagePixel i1 ~ FromFunctionPixel i2)
             => i1 -> i2
verticalFlip :: forall i1 i2.
(Image i1, FromFunction i2,
 ImagePixel i1 ~ FromFunctionPixel i2) =>
i1 -> i2
verticalFlip !i1
img =
    let line :: Int -> Int
line !Int
y' = Int
maxY forall a. Num a => a -> a -> a
- Int
y'
        {-# INLINE line #-}
        f :: Int -> Size -> ImagePixel i1
f !Int
y !(DIM0
Z :. Int
_ :. Int
x) = i1
img forall i. Image i => i -> Size -> ImagePixel i
! Int -> Int -> Size
ix2 Int
y Int
x
        {-# INLINE f #-}
    in forall i a.
FromFunction i =>
Size -> (Int -> a) -> (a -> Size -> FromFunctionPixel i) -> i
fromFunctionLine Size
size Int -> Int
line Int -> Size -> ImagePixel i1
f
  where
    !size :: Size
size@(DIM0
Z :. Int
h :. Int
_) = forall i. MaskedImage i => i -> Size
shape i1
img
    !maxY :: Int
maxY = Int
h forall a. Num a => a -> a -> a
- Int
1
{-# INLINABLE verticalFlip #-}

-- | Paints with a new value the pixels surrounding the given point of the image
-- which have the same value as the starting point.
floodFill :: (PrimMonad m, MutableImage i, Eq (ImagePixel (Freezed i)))
          => Point -> ImagePixel (Freezed i) -> i (PrimState m) -> m ()
floodFill :: forall (m :: * -> *) (i :: * -> *).
(PrimMonad m, MutableImage i, Eq (ImagePixel (Freezed i))) =>
Size -> ImagePixel (Freezed i) -> i (PrimState m) -> m ()
floodFill !Size
start !ImagePixel (Freezed i)
newVal !i (PrimState m)
img = do
    let !linearIX :: Int
linearIX = forall sh. Shape sh => sh -> sh -> Int
toLinearIndex Size
size Size
start
    ImagePixel (Freezed i)
val <- forall (i :: * -> *) (m :: * -> *).
(MutableImage i, PrimMonad m) =>
i (PrimState m) -> Int -> m (ImagePixel (Freezed i))
linearRead i (PrimState m)
img Int
linearIX
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ImagePixel (Freezed i)
val forall a. Eq a => a -> a -> Bool
/= ImagePixel (Freezed i)
newVal) forall a b. (a -> b) -> a -> b
$ -- No reason to repaint using the same color.
        ImagePixel (Freezed i) -> Size -> Int -> m ()
go ImagePixel (Freezed i)
val Size
start Int
linearIX
  where
    !size :: Size
size@(DIM0
Z :. Int
h :. Int
w) = forall (i :: * -> *) s. MutableImage i => i s -> Size
mShape i (PrimState m)
img

    -- Runs the flood-fill algorithm from the starting point then checks the
    -- pixels at the left and at the right of the point until their value
    -- change (scanLine). Then visits the upper and lower line of neighboring
    -- pixels (visitLine).

    go :: ImagePixel (Freezed i) -> Size -> Int -> m ()
go !ImagePixel (Freezed i)
val !(DIM0
Z :. Int
y :. Int
x) !Int
linearIX = do
        ImagePixel (Freezed i)
pix <- forall (i :: * -> *) (m :: * -> *).
(MutableImage i, PrimMonad m) =>
i (PrimState m) -> Int -> m (ImagePixel (Freezed i))
linearRead i (PrimState m)
img Int
linearIX

        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ImagePixel (Freezed i)
pix forall a. Eq a => a -> a -> Bool
== ImagePixel (Freezed i)
val) forall a b. (a -> b) -> a -> b
$ do
            let !minLineLinearIX :: Int
minLineLinearIX = Int
linearIX forall a. Num a => a -> a -> a
- Int
x
                !maxLineLinearIX :: Int
maxLineLinearIX = Int
minLineLinearIX forall a. Num a => a -> a -> a
+ Int
w forall a. Num a => a -> a -> a
- Int
1

            forall (i :: * -> *) (m :: * -> *).
(MutableImage i, PrimMonad m) =>
i (PrimState m) -> Int -> ImagePixel (Freezed i) -> m ()
linearWrite i (PrimState m)
img Int
linearIX ImagePixel (Freezed i)
newVal

            Int
stopLeft  <- ImagePixel (Freezed i)
-> (Int -> Bool) -> (Int -> Int) -> Int -> m Int
scanLine ImagePixel (Freezed i)
val (forall a. Ord a => a -> a -> Bool
< Int
minLineLinearIX) forall a. Enum a => a -> a
pred (Int
linearIX forall a. Num a => a -> a -> a
- Int
1)
            Int
stopRight <- ImagePixel (Freezed i)
-> (Int -> Bool) -> (Int -> Int) -> Int -> m Int
scanLine ImagePixel (Freezed i)
val (forall a. Ord a => a -> a -> Bool
> Int
maxLineLinearIX) forall a. Enum a => a -> a
succ (Int
linearIX forall a. Num a => a -> a -> a
+ Int
1)

            let !from :: Int
from  = Int
stopLeft  forall a. Num a => a -> a -> a
+ Int
1
                !to :: Int
to    = Int
stopRight forall a. Num a => a -> a -> a
- Int
1
                !xFrom :: Int
xFrom = Int
from forall a. Num a => a -> a -> a
- Int
minLineLinearIX

            forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
y forall a. Ord a => a -> a -> Bool
> Int
0) forall a b. (a -> b) -> a -> b
$
                ImagePixel (Freezed i) -> Int -> Size -> Int -> m ()
visitLine ImagePixel (Freezed i)
val (Int
to forall a. Num a => a -> a -> a
- Int
w) (Int -> Int -> Size
ix2 (Int
y forall a. Num a => a -> a -> a
- Int
1) Int
xFrom) (Int
from forall a. Num a => a -> a -> a
- Int
w)
            forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ((Int
y forall a. Num a => a -> a -> a
+ Int
1) forall a. Ord a => a -> a -> Bool
< Int
h) forall a b. (a -> b) -> a -> b
$
                ImagePixel (Freezed i) -> Int -> Size -> Int -> m ()
visitLine ImagePixel (Freezed i)
val (Int
to forall a. Num a => a -> a -> a
+ Int
w) (Int -> Int -> Size
ix2 (Int
y forall a. Num a => a -> a -> a
+ Int
1) Int
xFrom) (Int
from forall a. Num a => a -> a -> a
+ Int
w)

    scanLine :: ImagePixel (Freezed i)
-> (Int -> Bool) -> (Int -> Int) -> Int -> m Int
scanLine !ImagePixel (Freezed i)
val !Int -> Bool
stop !Int -> Int
next !Int
linearIX
        | Int -> Bool
stop Int
linearIX = forall (m :: * -> *) a. Monad m => a -> m a
return Int
linearIX
        | Bool
otherwise     = do
            ImagePixel (Freezed i)
pix <- forall (i :: * -> *) (m :: * -> *).
(MutableImage i, PrimMonad m) =>
i (PrimState m) -> Int -> m (ImagePixel (Freezed i))
linearRead i (PrimState m)
img Int
linearIX
            if ImagePixel (Freezed i)
pix forall a. Eq a => a -> a -> Bool
== ImagePixel (Freezed i)
val then do forall (i :: * -> *) (m :: * -> *).
(MutableImage i, PrimMonad m) =>
i (PrimState m) -> Int -> ImagePixel (Freezed i) -> m ()
linearWrite i (PrimState m)
img Int
linearIX ImagePixel (Freezed i)
newVal
                                  ImagePixel (Freezed i)
-> (Int -> Bool) -> (Int -> Int) -> Int -> m Int
scanLine ImagePixel (Freezed i)
val Int -> Bool
stop Int -> Int
next (Int -> Int
next Int
linearIX)
                          else forall (m :: * -> *) a. Monad m => a -> m a
return Int
linearIX

    visitLine :: ImagePixel (Freezed i) -> Int -> Size -> Int -> m ()
visitLine !ImagePixel (Freezed i)
val !Int
maxLinearIX !pt :: Size
pt@(DIM0 :. Int
y :. Int
x) !Int
linearIX
        | Int
linearIX forall a. Ord a => a -> a -> Bool
> Int
maxLinearIX = forall (m :: * -> *) a. Monad m => a -> m a
return ()
        | Bool
otherwise              = do
            ImagePixel (Freezed i) -> Size -> Int -> m ()
go ImagePixel (Freezed i)
val Size
pt Int
linearIX
            ImagePixel (Freezed i) -> Int -> Size -> Int -> m ()
visitLine ImagePixel (Freezed i)
val Int
maxLinearIX (DIM0 :. Int
y forall tail head. tail -> head -> tail :. head
:. (Int
x forall a. Num a => a -> a -> a
+ Int
1)) (Int
linearIX forall a. Num a => a -> a -> a
+ Int
1)
{-# INLINABLE floodFill #-}

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

ratio :: Integral a => a -> RatioInt
ratio :: forall a. Integral a => a -> RatioInt
ratio = forall a b. (Integral a, Num b) => a -> b
fromIntegral