BMP.hs

{-# LANGUAGE PackageImports, PatternGuards, ExplicitForAll  #-} 

-- | Reading and writing arrays as uncompressed 24 or 32 bit Windows BMP files.
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

-- NOTE: We set most of these functions as NOINLINE so it's easier to understand
--       what's going on in the core programs. The top-level IO functions are
--       only called a few times each, so it doesn't matter if they're not
--       worker/wrappered etc.
        
-- Read -------------------------------------------------------------------------------------------
-- | Read RGB components from a BMP file.
readImageFromBMP
        :: FilePath
        -> IO (Either Error (Array U DIM2 (Word8, Word8, Word8)))

{-# NOINLINE readImageFromBMP #-}
readImageFromBMP filePath
 = do   ebmp    <- readBMP filePath

        case ebmp of
         Left err       -> return $ Left err
         Right bmp      
          -> do arr     <- readImageFromBMP' bmp
                return  $ Right arr

readImageFromBMP' bmp
 = do   let (width, height) = bmpDimensions bmp

        let arr         = R.fromByteString (Z :. height :. width * 4)
                        $ unpackBMPToRGBA32 bmp

        let shapeFn _   = Z :. height :. width

        -- Read out the components into their own arrays, 
        -- skipping the alpha channel.
        vecRed         <- computeP 
                        $ unsafeTraverse arr shapeFn
                               (\get (sh :. x) -> get (sh :. (x * 4)))

        vecGreen       <- computeP 
                        $ unsafeTraverse arr shapeFn
                                (\get (sh :. x) -> get (sh :. (x * 4 + 1)))

        vecBlue        <- computeP
                        $ unsafeTraverse arr shapeFn
                                (\get (sh :. x) -> get (sh :. (x * 4 + 2)))

        -- O(1). zip the components together
        let vecRGB     = U.zip3 (toUnboxed vecRed)
                                (toUnboxed vecGreen)
                                (toUnboxed vecBlue)
        
        return $ fromUnboxed (Z :. height :. width) vecRGB



-- | Write RGB components to a BMP file.
writeImageToBMP
        :: FilePath
        -> Array U DIM2 (Word8, Word8, Word8)
        -> IO ()

{-# NOINLINE writeImageToBMP #-}
writeImageToBMP fileName arrRGB
 = do   let sh@(Z :. height :. width)
                        = extent arrRGB

        -- O(1). unzip the components
        let (vecRed, vecGreen, vecBlue)
                = U.unzip3 $ toUnboxed arrRGB


        -- Create a bytestring with all the data
        ptr     <- mallocBytes (height * width * 4)
        fptr    <- newForeignPtr finalizerFree ptr

        computeIntoP fptr 
         $ interleave4 
                (fromUnboxed sh vecRed)
                (fromUnboxed sh vecGreen)
                (fromUnboxed sh vecBlue)
                (fromFunction sh (\_ -> 255))

        -- Pack the data into a BMP file and write it out.
        withForeignPtr fptr
         $ \ptr' -> do
                bs      <- unsafePackCStringFinalizer ptr' (width * height * 4) (return ())
                let bmp = packRGBA32ToBMP width height bs
                writeBMP fileName bmp


{-
-- Normalise --------------------------------------------------------------------------------------
-- | Normalise a matrix to to [0 .. 1], discarding negative values.
--      If the maximum value is 0 then return the array unchanged.
normalisePositive01
        :: (Shape sh, Fractional a, Ord a)
        => Array sh a
        -> Array sh a

{-# INLINE normalisePositive01 #-}
normalisePositive01 arr 
 = let  mx              = foldAll max 0 arr
        elemFn x
         | x >= 0       = x / mx
         | otherwise    = x
   in   mx `seq`
         if mx == 0 
          then arr
          else R.map elemFn arr

-}