{-# LANGUAGE PackageImports, PatternGuards, ExplicitForAll #-}
module Data.Array.Repa.IO.BMP
( readImageFromBMP
, writeImageToBMP)
where
import Data.Array.Repa as R
import Data.Array.Repa.Unsafe as R
import Data.Array.Repa.Repr.ForeignPtr as R
import Data.Array.Repa.Repr.ByteString as R
import Data.Vector.Unboxed as U
import Prelude as P
import Foreign.ForeignPtr
import Foreign.Marshal.Alloc
import Data.ByteString.Unsafe as B
import Codec.BMP
import Data.Word
readImageFromBMP
:: FilePath
-> IO (Either Error (Array U DIM2 (Word8, Word8, Word8)))
{-# NOINLINE readImageFromBMP #-}
readImageFromBMP :: FilePath -> IO (Either Error (Array U DIM2 (Word8, Word8, Word8)))
readImageFromBMP FilePath
filePath
= do Either Error BMP
ebmp <- FilePath -> IO (Either Error BMP)
readBMP FilePath
filePath
case Either Error BMP
ebmp of
Left Error
err -> Either Error (Array U DIM2 (Word8, Word8, Word8))
-> IO (Either Error (Array U DIM2 (Word8, Word8, Word8)))
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Error (Array U DIM2 (Word8, Word8, Word8))
-> IO (Either Error (Array U DIM2 (Word8, Word8, Word8))))
-> Either Error (Array U DIM2 (Word8, Word8, Word8))
-> IO (Either Error (Array U DIM2 (Word8, Word8, Word8)))
forall a b. (a -> b) -> a -> b
$ Error -> Either Error (Array U DIM2 (Word8, Word8, Word8))
forall a b. a -> Either a b
Left Error
err
Right BMP
bmp
-> do Array U DIM2 (Word8, Word8, Word8)
arr <- BMP -> IO (Array U DIM2 (Word8, Word8, Word8))
forall (m :: * -> *).
Monad m =>
BMP -> m (Array U DIM2 (Word8, Word8, Word8))
readImageFromBMP' BMP
bmp
Either Error (Array U DIM2 (Word8, Word8, Word8))
-> IO (Either Error (Array U DIM2 (Word8, Word8, Word8)))
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Error (Array U DIM2 (Word8, Word8, Word8))
-> IO (Either Error (Array U DIM2 (Word8, Word8, Word8))))
-> Either Error (Array U DIM2 (Word8, Word8, Word8))
-> IO (Either Error (Array U DIM2 (Word8, Word8, Word8)))
forall a b. (a -> b) -> a -> b
$ Array U DIM2 (Word8, Word8, Word8)
-> Either Error (Array U DIM2 (Word8, Word8, Word8))
forall a b. b -> Either a b
Right Array U DIM2 (Word8, Word8, Word8)
arr
readImageFromBMP' :: BMP -> m (Array U DIM2 (Word8, Word8, Word8))
readImageFromBMP' BMP
bmp
= do let (Int
width, Int
height) = BMP -> (Int, Int)
bmpDimensions BMP
bmp
let arr :: Array B DIM2 Word8
arr = DIM2 -> ByteString -> Array B DIM2 Word8
forall sh. sh -> ByteString -> Array B sh Word8
R.fromByteString (Z
Z Z -> Int -> Z :. Int
forall tail head. tail -> head -> tail :. head
:. Int
height (Z :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4)
(ByteString -> Array B DIM2 Word8)
-> ByteString -> Array B DIM2 Word8
forall a b. (a -> b) -> a -> b
$ BMP -> ByteString
unpackBMPToRGBA32 BMP
bmp
let shapeFn :: p -> DIM2
shapeFn p
_ = Z
Z Z -> Int -> Z :. Int
forall tail head. tail -> head -> tail :. head
:. Int
height (Z :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
width
Array U DIM2 Word8
vecRed <- Array D DIM2 Word8 -> m (Array U DIM2 Word8)
forall r1 sh e r2 (m :: * -> *).
(Load r1 sh e, Target r2 e, Source r2 e, Monad m) =>
Array r1 sh e -> m (Array r2 sh e)
computeP
(Array D DIM2 Word8 -> m (Array U DIM2 Word8))
-> Array D DIM2 Word8 -> m (Array U DIM2 Word8)
forall a b. (a -> b) -> a -> b
$ Array B DIM2 Word8
-> (DIM2 -> DIM2)
-> ((DIM2 -> Word8) -> DIM2 -> Word8)
-> Array D DIM2 Word8
forall r sh sh' a b.
(Source r a, Shape sh) =>
Array r sh a
-> (sh -> sh') -> ((sh -> a) -> sh' -> b) -> Array D sh' b
unsafeTraverse Array B DIM2 Word8
arr DIM2 -> DIM2
forall p. p -> DIM2
shapeFn
(\DIM2 -> Word8
get (Z :. Int
sh :. Int
x) -> DIM2 -> Word8
get (Z :. Int
sh (Z :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4)))
Array U DIM2 Word8
vecGreen <- Array D DIM2 Word8 -> m (Array U DIM2 Word8)
forall r1 sh e r2 (m :: * -> *).
(Load r1 sh e, Target r2 e, Source r2 e, Monad m) =>
Array r1 sh e -> m (Array r2 sh e)
computeP
(Array D DIM2 Word8 -> m (Array U DIM2 Word8))
-> Array D DIM2 Word8 -> m (Array U DIM2 Word8)
forall a b. (a -> b) -> a -> b
$ Array B DIM2 Word8
-> (DIM2 -> DIM2)
-> ((DIM2 -> Word8) -> DIM2 -> Word8)
-> Array D DIM2 Word8
forall r sh sh' a b.
(Source r a, Shape sh) =>
Array r sh a
-> (sh -> sh') -> ((sh -> a) -> sh' -> b) -> Array D sh' b
unsafeTraverse Array B DIM2 Word8
arr DIM2 -> DIM2
forall p. p -> DIM2
shapeFn
(\DIM2 -> Word8
get (Z :. Int
sh :. Int
x) -> DIM2 -> Word8
get (Z :. Int
sh (Z :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)))
Array U DIM2 Word8
vecBlue <- Array D DIM2 Word8 -> m (Array U DIM2 Word8)
forall r1 sh e r2 (m :: * -> *).
(Load r1 sh e, Target r2 e, Source r2 e, Monad m) =>
Array r1 sh e -> m (Array r2 sh e)
computeP
(Array D DIM2 Word8 -> m (Array U DIM2 Word8))
-> Array D DIM2 Word8 -> m (Array U DIM2 Word8)
forall a b. (a -> b) -> a -> b
$ Array B DIM2 Word8
-> (DIM2 -> DIM2)
-> ((DIM2 -> Word8) -> DIM2 -> Word8)
-> Array D DIM2 Word8
forall r sh sh' a b.
(Source r a, Shape sh) =>
Array r sh a
-> (sh -> sh') -> ((sh -> a) -> sh' -> b) -> Array D sh' b
unsafeTraverse Array B DIM2 Word8
arr DIM2 -> DIM2
forall p. p -> DIM2
shapeFn
(\DIM2 -> Word8
get (Z :. Int
sh :. Int
x) -> DIM2 -> Word8
get (Z :. Int
sh (Z :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)))
let vecRGB :: Vector (Word8, Word8, Word8)
vecRGB = Vector Word8
-> Vector Word8 -> Vector Word8 -> Vector (Word8, Word8, Word8)
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
Vector a -> Vector b -> Vector c -> Vector (a, b, c)
U.zip3 (Array U DIM2 Word8 -> Vector Word8
forall sh e. Array U sh e -> Vector e
toUnboxed Array U DIM2 Word8
vecRed)
(Array U DIM2 Word8 -> Vector Word8
forall sh e. Array U sh e -> Vector e
toUnboxed Array U DIM2 Word8
vecGreen)
(Array U DIM2 Word8 -> Vector Word8
forall sh e. Array U sh e -> Vector e
toUnboxed Array U DIM2 Word8
vecBlue)
Array U DIM2 (Word8, Word8, Word8)
-> m (Array U DIM2 (Word8, Word8, Word8))
forall (m :: * -> *) a. Monad m => a -> m a
return (Array U DIM2 (Word8, Word8, Word8)
-> m (Array U DIM2 (Word8, Word8, Word8)))
-> Array U DIM2 (Word8, Word8, Word8)
-> m (Array U DIM2 (Word8, Word8, Word8))
forall a b. (a -> b) -> a -> b
$ DIM2
-> Vector (Word8, Word8, Word8)
-> Array U DIM2 (Word8, Word8, Word8)
forall sh e. sh -> Vector e -> Array U sh e
fromUnboxed (Z
Z Z -> Int -> Z :. Int
forall tail head. tail -> head -> tail :. head
:. Int
height (Z :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
width) Vector (Word8, Word8, Word8)
vecRGB
writeImageToBMP
:: FilePath
-> Array U DIM2 (Word8, Word8, Word8)
-> IO ()
{-# NOINLINE writeImageToBMP #-}
writeImageToBMP :: FilePath -> Array U DIM2 (Word8, Word8, Word8) -> IO ()
writeImageToBMP FilePath
fileName Array U DIM2 (Word8, Word8, Word8)
arrRGB
= do let sh :: DIM2
sh@(Z
Z :. Int
height :. Int
width)
= Array U DIM2 (Word8, Word8, Word8) -> DIM2
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array U DIM2 (Word8, Word8, Word8)
arrRGB
let (Vector Word8
vecRed, Vector Word8
vecGreen, Vector Word8
vecBlue)
= Vector (Word8, Word8, Word8)
-> (Vector Word8, Vector Word8, Vector Word8)
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
Vector (a, b, c) -> (Vector a, Vector b, Vector c)
U.unzip3 (Vector (Word8, Word8, Word8)
-> (Vector Word8, Vector Word8, Vector Word8))
-> Vector (Word8, Word8, Word8)
-> (Vector Word8, Vector Word8, Vector Word8)
forall a b. (a -> b) -> a -> b
$ Array U DIM2 (Word8, Word8, Word8) -> Vector (Word8, Word8, Word8)
forall sh e. Array U sh e -> Vector e
toUnboxed Array U DIM2 (Word8, Word8, Word8)
arrRGB
Ptr Word8
ptr <- Int -> IO (Ptr Word8)
forall a. Int -> IO (Ptr a)
mallocBytes (Int
height Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4)
ForeignPtr Word8
fptr <- FinalizerPtr Word8 -> Ptr Word8 -> IO (ForeignPtr Word8)
forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
newForeignPtr FinalizerPtr Word8
forall a. FinalizerPtr a
finalizerFree Ptr Word8
ptr
ForeignPtr Word8 -> Array D DIM2 Word8 -> IO ()
forall r1 sh e.
(Load r1 sh e, Storable e) =>
ForeignPtr e -> Array r1 sh e -> IO ()
computeIntoP ForeignPtr Word8
fptr
(Array D DIM2 Word8 -> IO ()) -> Array D DIM2 Word8 -> IO ()
forall a b. (a -> b) -> a -> b
$ Array U DIM2 Word8
-> Array U DIM2 Word8
-> Array U DIM2 Word8
-> Array D DIM2 Word8
-> Array D DIM2 Word8
forall sh r1 a r2 r3 r4.
(Shape sh, Source r1 a, Source r2 a, Source r3 a, Source r4 a) =>
Array r1 (sh :. Int) a
-> Array r2 (sh :. Int) a
-> Array r3 (sh :. Int) a
-> Array r4 (sh :. Int) a
-> Array D (sh :. Int) a
interleave4
(DIM2 -> Vector Word8 -> Array U DIM2 Word8
forall sh e. sh -> Vector e -> Array U sh e
fromUnboxed DIM2
sh Vector Word8
vecRed)
(DIM2 -> Vector Word8 -> Array U DIM2 Word8
forall sh e. sh -> Vector e -> Array U sh e
fromUnboxed DIM2
sh Vector Word8
vecGreen)
(DIM2 -> Vector Word8 -> Array U DIM2 Word8
forall sh e. sh -> Vector e -> Array U sh e
fromUnboxed DIM2
sh Vector Word8
vecBlue)
(DIM2 -> (DIM2 -> Word8) -> Array D DIM2 Word8
forall sh a. sh -> (sh -> a) -> Array D sh a
fromFunction DIM2
sh (\DIM2
_ -> Word8
255))
ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fptr
((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr' -> do
ByteString
bs <- Ptr Word8 -> Int -> IO () -> IO ByteString
unsafePackCStringFinalizer Ptr Word8
ptr' (Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
height Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4) (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
let bmp :: BMP
bmp = Int -> Int -> ByteString -> BMP
packRGBA32ToBMP Int
width Int
height ByteString
bs
FilePath -> BMP -> IO ()
writeBMP FilePath
fileName BMP
bmp