{-# LANGUAGE DeriveAnyClass, DeriveGeneric, FlexibleContexts, OverloadedStrings #-}
module Graphics.OpenEXR
( Image(..)
, PixelRGBF(..)
, CompressionType(..)
, writeFile
) where
import Prelude hiding (writeFile)
import Codec.Compression.Zlib (compress)
import Control.DeepSeq (NFData)
import Data.Binary (Binary, get, put)
import Data.Binary.IEEE754 (putFloat32le)
import Data.Binary.Put (Put, putLazyByteString, putWord8, putWord32le, putWord64le, runPut)
import qualified Data.ByteString.Lazy as BL
import Data.List.Split (chunksOf)
import qualified Data.Vector as V
import qualified Data.Vector.Split as V
import Data.Word (Word8, Word32)
import GHC.Generics (Generic)
data Image = ImageRGBF
{ imageWidth :: !Int
, imageHeight :: !Int
, imageData :: V.Vector PixelRGBF
} deriving (Generic, NFData, Show)
data PixelRGBF = PixelRGBF !Float !Float !Float
deriving (Generic, NFData, Show)
data CompressionType = NoCompression
| ZipsCompression
| ZipCompression
data Attribute = Attribute BL.ByteString AttributeType
data AttributeType = Box2i (Word32, Word32, Word32, Word32)
| Chlist [Channel]
| Compression CompressionType
| Float Float
| LineOrder LineOrderType
| V2f (Float, Float)
data Channel = Channel
{ name :: BL.ByteString
, pixelType :: PixelType
, pLinear :: Word8
, reserved :: (Word8, Word8, Word8)
, xSampling :: Word32
, ySampling :: Word32
}
data PixelType = PixelUint | PixelHalf | PixelFloat
deriving Enum
data LineOrderType = IncreasingY | DecreasingY | RandomY
deriving Enum
instance Binary Attribute where
get = undefined
put (Attribute n t) = putLazyByteString (BL.snoc n 0) >> put t
instance Binary AttributeType where
get = undefined
put a@(Box2i (x, y, z, w)) = putLazyByteString "box2i\0" >> putWord32le (attrSize a) >> putWord32le x >> putWord32le y >> putWord32le z >> putWord32le w
put a@(Chlist xs) = putLazyByteString "chlist\0" >> putWord32le (attrSize a) >> mapM_ put xs >> put '\0'
put a@(Compression x) = putLazyByteString "compression\0" >> putWord32le (attrSize a) >> put x
put a@(Float x) = putLazyByteString "float\0" >> putWord32le (attrSize a) >> putFloat32le x
put a@(LineOrder x) = putLazyByteString "lineOrder\0" >> putWord32le (attrSize a) >> put x
put a@(V2f (x, y)) = putLazyByteString "v2f\0" >> putWord32le (attrSize a) >> putFloat32le x >> putFloat32le y
instance Binary Channel where
get = undefined
put x = do
putLazyByteString (BL.snoc (name x) 0)
put (pixelType x)
put (pLinear x)
put (reserved x)
putWord32le (xSampling x)
putWord32le (ySampling x)
instance Binary PixelType where
get = undefined
put = putWord32le . fromIntegral . fromEnum
instance Enum CompressionType where
fromEnum NoCompression = 0
fromEnum ZipsCompression = 2
fromEnum ZipCompression = 3
toEnum _ = undefined
instance Binary CompressionType where
get = undefined
put = putWord8 . fromIntegral . fromEnum
instance Binary LineOrderType where
get = undefined
put = putWord8 . fromIntegral . fromEnum
writeFile :: FilePath -> Image -> CompressionType -> IO ()
writeFile filepath img compression = BL.writeFile filepath (header `BL.append` offTable `BL.append` (BL.concat cs))
where header = runPut (magicNumber >> versionField >> mapM_ put attributes >> put '\0')
offTable = runPut (offsetTable (fromIntegral $ BL.length header) ((fromIntegral . BL.length) <$> cs))
cs = runPut <$> chunks img compression
attributes = [ Attribute "channels" channels
, Attribute "compression" (Compression compression)
, Attribute "dataWindow" (Box2i (0, 0, w, h))
, Attribute "displayWindow" (Box2i (0, 0, w, h))
, Attribute "lineOrder" (LineOrder IncreasingY)
, Attribute "pixelAspectRatio" (Float 1)
, Attribute "screenWindowCenter" (V2f (0, 0))
, Attribute "screenWindowWidth" (Float 1)]
channels = Chlist [ Channel "B" PixelFloat 0 (0, 0, 0) 1 1
, Channel "G" PixelFloat 0 (0, 0, 0) 1 1
, Channel "R" PixelFloat 0 (0, 0, 0) 1 1
]
w = fromIntegral (imageWidth img) - 1
h = fromIntegral (imageHeight img) - 1
attrSize :: AttributeType -> Word32
attrSize (Box2i _) = 16
attrSize (Chlist xs) = fromIntegral ((sum . map (\x -> (BL.length . name $ x) + 1 + 4 + 1 + 3 + 4 + 4) $ xs) + 1)
attrSize (Compression _) = 1
attrSize (Float _) = 4
attrSize (LineOrder _) = 1
attrSize (V2f _) = 8
magicNumber :: Put
magicNumber = mapM_ putWord8 [0x76, 0x2F, 0x31, 0x01]
versionField :: Put
versionField = mapM_ putWord8 [0x02, 0x00, 0x00, 0x00]
offsetTable :: Int -> [Int] -> Put
offsetTable offset chunksLengths = mapM_ (\x -> putWord64le (fromIntegral (offset + 8*nChunks + x))) . scanl (+) 0 . init $ chunksLengths
where nChunks = length chunksLengths
scanlinesPerChunk :: Num a => CompressionType -> a
scanlinesPerChunk ZipCompression = 16
scanlinesPerChunk _ = 1
transform :: CompressionType -> (BL.ByteString -> BL.ByteString)
transform ZipsCompression = compress . preZip
transform ZipCompression = compress . preZip
transform _ = id
chunks :: Image -> CompressionType -> [Put]
chunks (ImageRGBF width _ px) c = uncurry f <$> zip [0,yStride..] (chunksOf yStride cs)
where cs = V.chunksOf width px
yStride = scanlinesPerChunk c
f y xs = do
let d = runPut . mapM_ scanline $ xs
transformed = transform c d
result = if BL.length d <= BL.length transformed then d else transformed
putWord32le (fromIntegral y)
putWord32le (fromIntegral (BL.length result))
putLazyByteString result
scanline :: V.Vector PixelRGBF -> Put
scanline xs = do
mapM_ (\(PixelRGBF _ _ b) -> putFloat32le b) xs
mapM_ (\(PixelRGBF _ g _) -> putFloat32le g) xs
mapM_ (\(PixelRGBF r _ _) -> putFloat32le r) xs
preZip :: BL.ByteString -> BL.ByteString
preZip d = BL.pack (x:zipWith predictor (x:xs ++ ys) (xs ++ ys))
where (x:xs, ys) = deinterleave (BL.unpack d)
predictor p value = fromIntegral value - fromIntegral p + (128 + 256)
deinterleave :: [a] -> ([a], [a])
deinterleave [] = ([], [])
deinterleave [x] = ([x], [])
deinterleave [x,y] = ([x], [y])
deinterleave (x:y:xs) = let (as, bs) = deinterleave xs in (x:as, y:bs)