1{-# LANGUAGE CPP #-} 2{-# LANGUAGE TypeFamilies #-} 3{-# LANGUAGE TupleSections #-} 4-- | Module dedicated of Radiance file decompression (.hdr or .pic) file. 5-- Radiance file format is used for High dynamic range imaging. 6module Codec.Picture.HDR( decodeHDR 7 , decodeHDRWithMetadata 8 , encodeHDR 9 , encodeRawHDR 10 , encodeRLENewStyleHDR 11 , writeHDR 12 , writeRLENewStyleHDR 13 ) where 14 15#if !MIN_VERSION_base(4,8,0) 16import Control.Applicative( pure, (<*>), (<$>) ) 17#endif 18 19import Data.Bits( Bits, (.&.), (.|.), unsafeShiftL, unsafeShiftR ) 20import Data.Char( ord, chr, isDigit ) 21import Data.Word( Word8 ) 22 23#if !MIN_VERSION_base(4,11,0) 24import Data.Monoid( (<>) ) 25#endif 26 27import Control.Monad( when, foldM, foldM_, forM, forM_, unless ) 28import Control.Monad.Trans.Class( lift ) 29import qualified Data.ByteString as B 30import qualified Data.ByteString.Lazy as L 31import qualified Data.ByteString.Char8 as BC 32 33import Data.List( partition ) 34import Data.Binary( Binary( .. ), encode ) 35import Data.Binary.Get( Get, getByteString, getWord8 ) 36import Data.Binary.Put( putByteString, putLazyByteString ) 37 38import Control.Monad.ST( ST, runST ) 39import Foreign.Storable ( Storable ) 40import Control.Monad.Primitive ( PrimState, PrimMonad ) 41import qualified Data.Vector.Storable as V 42import qualified Data.Vector.Storable.Mutable as M 43 44import Codec.Picture.Metadata( Metadatas 45 , SourceFormat( SourceHDR ) 46 , basicMetadata ) 47import Codec.Picture.InternalHelper 48import Codec.Picture.Types 49import Codec.Picture.VectorByteConversion 50 51#if MIN_VERSION_transformers(0, 4, 0) 52import Control.Monad.Trans.Except( ExceptT, throwE, runExceptT ) 53#else 54-- Transfomers 0.3 compat 55import Control.Monad.Trans.Error( Error, ErrorT, throwError, runErrorT ) 56 57type ExceptT = ErrorT 58 59throwE :: (Monad m, Error e) => e -> ErrorT e m a 60throwE = throwError 61 62runExceptT :: ErrorT e m a -> m (Either e a) 63runExceptT = runErrorT 64#endif 65 66{-# INLINE (.<<.) #-} 67(.<<.), (.>>.) :: (Bits a) => a -> Int -> a 68(.<<.) = unsafeShiftL 69(.>>.) = unsafeShiftR 70 71{-# INLINE (.<-.) #-} 72(.<-.) :: (PrimMonad m, Storable a) 73 => M.STVector (PrimState m) a -> Int -> a -> m () 74(.<-.) = M.write 75 {-M.unsafeWrite-} 76 77type HDRReader s a = ExceptT String (ST s) a 78 79data RGBE = RGBE !Word8 !Word8 !Word8 !Word8 80 81instance Binary RGBE where 82 put (RGBE r g b e) = put r >> put g >> put b >> put e 83 get = RGBE <$> get <*> get <*> get <*> get 84 85checkLineLength :: RGBE -> Int 86checkLineLength (RGBE _ _ a b) = 87 (fromIntegral a .<<. 8) .|. fromIntegral b 88 89isNewRunLengthMarker :: RGBE -> Bool 90isNewRunLengthMarker (RGBE 2 2 _ _) = True 91isNewRunLengthMarker _ = False 92 93data RadianceFormat = 94 FormatRGBE 95 | FormatXYZE 96 97radiance32bitRleRGBEFormat, radiance32bitRleXYZEFromat :: B.ByteString 98radiance32bitRleRGBEFormat = BC.pack "32-bit_rle_rgbe" 99radiance32bitRleXYZEFromat = BC.pack "32-bit_rle_xyze" 100 101instance Binary RadianceFormat where 102 put FormatRGBE = putByteString radiance32bitRleRGBEFormat 103 put FormatXYZE = putByteString radiance32bitRleXYZEFromat 104 105 get = getByteString (B.length radiance32bitRleRGBEFormat) >>= format 106 where format sig 107 | sig == radiance32bitRleRGBEFormat = pure FormatRGBE 108 | sig == radiance32bitRleXYZEFromat = pure FormatXYZE 109 | otherwise = fail "Unrecognized Radiance format" 110 111toRGBE :: PixelRGBF -> RGBE 112toRGBE (PixelRGBF r g b) 113 | d <= 1e-32 = RGBE 0 0 0 0 114 | otherwise = RGBE (fix r) (fix g) (fix b) (fromIntegral $ e + 128) 115 where d = maximum [r, g, b] 116 e = exponent d 117 coeff = significand d * 255.9999 / d 118 fix v = truncate $ v * coeff 119 120 121dropUntil :: Word8 -> Get () 122dropUntil c = getWord8 >>= inner 123 where inner val | val == c = pure () 124 inner _ = getWord8 >>= inner 125 126getUntil :: (Word8 -> Bool) -> B.ByteString -> Get B.ByteString 127getUntil f initialAcc = getWord8 >>= inner initialAcc 128 where inner acc c | f c = pure acc 129 inner acc c = getWord8 >>= inner (B.snoc acc c) 130 131data RadianceHeader = RadianceHeader 132 { radianceInfos :: [(B.ByteString, B.ByteString)] 133 , radianceFormat :: RadianceFormat 134 , radianceHeight :: !Int 135 , radianceWidth :: !Int 136 , radianceData :: L.ByteString 137 } 138 139radianceFileSignature :: B.ByteString 140radianceFileSignature = BC.pack "#?RADIANCE\n" 141 142unpackColor :: L.ByteString -> Int -> RGBE 143unpackColor str idx = RGBE (at 0) (at 1) (at 2) (at 3) 144 where at n = L.index str . fromIntegral $ idx + n 145 146storeColor :: M.STVector s Word8 -> Int -> RGBE -> ST s () 147storeColor vec idx (RGBE r g b e) = do 148 (vec .<-. (idx + 0)) r 149 (vec .<-. (idx + 1)) g 150 (vec .<-. (idx + 2)) b 151 (vec .<-. (idx + 3)) e 152 153parsePair :: Char -> Get (B.ByteString, B.ByteString) 154parsePair firstChar = do 155 let eol c = c == fromIntegral (ord '\n') 156 line <- getUntil eol B.empty 157 case BC.split '=' line of 158 [] -> pure (BC.singleton firstChar, B.empty) 159 [val] -> pure (BC.singleton firstChar, val) 160 [key, val] -> pure (BC.singleton firstChar <> key, val) 161 (key : vals) -> pure (BC.singleton firstChar <> key, B.concat vals) 162 163decodeInfos :: Get [(B.ByteString, B.ByteString)] 164decodeInfos = do 165 char <- getChar8 166 case char of 167 -- comment 168 '#' -> dropUntil (fromIntegral $ ord '\n') >> decodeInfos 169 -- end of header, no more information 170 '\n' -> pure [] 171 -- Classical parsing 172 c -> (:) <$> parsePair c <*> decodeInfos 173 174 175-- | Decode an HDR (radiance) image, the resulting image can be: 176-- 177-- * 'ImageRGBF' 178-- 179decodeHDR :: B.ByteString -> Either String DynamicImage 180decodeHDR = fmap fst . decodeHDRWithMetadata 181 182-- | Equivalent to decodeHDR but with aditional metadatas. 183decodeHDRWithMetadata :: B.ByteString -> Either String (DynamicImage, Metadatas) 184decodeHDRWithMetadata str = runST $ runExceptT $ 185 case runGet decodeHeader $ L.fromChunks [str] of 186 Left err -> throwE err 187 Right rez -> 188 let meta = basicMetadata SourceHDR (abs $ radianceWidth rez) (abs $ radianceHeight rez) in 189 (, meta) . ImageRGBF <$> (decodeRadiancePicture rez >>= lift . unsafeFreezeImage) 190 191getChar8 :: Get Char 192getChar8 = chr . fromIntegral <$> getWord8 193 194isSign :: Char -> Bool 195isSign c = c == '+' || c == '-' 196 197isAxisLetter :: Char -> Bool 198isAxisLetter c = c == 'X' || c == 'Y' 199 200decodeNum :: Get Int 201decodeNum = do 202 sign <- getChar8 203 letter <- getChar8 204 space <- getChar8 205 206 unless (isSign sign && isAxisLetter letter && space == ' ') 207 (fail "Invalid radiance size declaration") 208 209 let numDec acc c | isDigit c = 210 getChar8 >>= numDec (acc * 10 + ord c - ord '0') 211 numDec acc _ 212 | sign == '-' = pure $ negate acc 213 | otherwise = pure acc 214 215 getChar8 >>= numDec 0 216 217copyPrevColor :: M.STVector s Word8 -> Int -> ST s () 218copyPrevColor scanLine idx = do 219 r <- scanLine `M.unsafeRead` (idx - 4) 220 g <- scanLine `M.unsafeRead` (idx - 3) 221 b <- scanLine `M.unsafeRead` (idx - 2) 222 e <- scanLine `M.unsafeRead` (idx - 1) 223 224 (scanLine `M.unsafeWrite` (idx + 0)) r 225 (scanLine `M.unsafeWrite` (idx + 1)) g 226 (scanLine `M.unsafeWrite` (idx + 2)) b 227 (scanLine `M.unsafeWrite` (idx + 3)) e 228 229oldStyleRLE :: L.ByteString -> Int -> M.STVector s Word8 230 -> HDRReader s Int 231oldStyleRLE inputData initialIdx scanLine = inner initialIdx 0 0 232 where maxOutput = M.length scanLine 233 maxInput = fromIntegral $ L.length inputData 234 235 inner readIdx writeIdx _ 236 | readIdx >= maxInput || writeIdx >= maxOutput = pure readIdx 237 inner readIdx writeIdx shift = do 238 let color@(RGBE r g b e) = unpackColor inputData readIdx 239 isRun = r == 1 && g == 1 && b == 1 240 241 if not isRun 242 then do 243 lift $ storeColor scanLine writeIdx color 244 inner (readIdx + 4) (writeIdx + 4) 0 245 246 else do 247 let count = fromIntegral e .<<. shift 248 lift $ forM_ [0 .. count] $ \i -> copyPrevColor scanLine (writeIdx + 4 * i) 249 inner (readIdx + 4) (writeIdx + 4 * count) (shift + 8) 250 251newStyleRLE :: L.ByteString -> Int -> M.STVector s Word8 252 -> HDRReader s Int 253newStyleRLE inputData initialIdx scanline = foldM inner initialIdx [0 .. 3] 254 where dataAt idx 255 | fromIntegral idx >= maxInput = throwE $ "Read index out of bound (" ++ show idx ++ ")" 256 | otherwise = pure $ L.index inputData (fromIntegral idx) 257 258 maxOutput = M.length scanline 259 maxInput = fromIntegral $ L.length inputData 260 stride = 4 261 262 263 strideSet count destIndex _ | endIndex > maxOutput + stride = 264 throwE $ "Out of bound HDR scanline " ++ show endIndex ++ " (max " ++ show maxOutput ++ ")" 265 where endIndex = destIndex + count * stride 266 strideSet count destIndex val = aux destIndex count 267 where aux i 0 = pure i 268 aux i c = do 269 lift $ (scanline .<-. i) val 270 aux (i + stride) (c - 1) 271 272 273 strideCopy _ count destIndex 274 | writeEndBound > maxOutput + stride = throwE "Out of bound HDR scanline" 275 where writeEndBound = destIndex + count * stride 276 strideCopy sourceIndex count destIndex = aux sourceIndex destIndex count 277 where aux _ j 0 = pure j 278 aux i j c = do 279 val <- dataAt i 280 lift $ (scanline .<-. j) val 281 aux (i + 1) (j + stride) (c - 1) 282 283 inner readIdx writeIdx 284 | readIdx >= maxInput || writeIdx >= maxOutput = pure readIdx 285 inner readIdx writeIdx = do 286 code <- dataAt readIdx 287 if code > 128 288 then do 289 let repeatCount = fromIntegral code .&. 0x7F 290 newVal <- dataAt $ readIdx + 1 291 endIndex <- strideSet repeatCount writeIdx newVal 292 inner (readIdx + 2) endIndex 293 294 else do 295 let iCode = fromIntegral code 296 endIndex <- strideCopy (readIdx + 1) iCode writeIdx 297 inner (readIdx + iCode + 1) endIndex 298 299instance Binary RadianceHeader where 300 get = decodeHeader 301 put hdr = do 302 putByteString radianceFileSignature 303 putByteString $ BC.pack "FORMAT=" 304 put $ radianceFormat hdr 305 let sizeString = 306 BC.pack $ "\n\n-Y " ++ show (radianceHeight hdr) 307 ++ " +X " ++ show (radianceWidth hdr) ++ "\n" 308 putByteString sizeString 309 putLazyByteString $ radianceData hdr 310 311 312decodeHeader :: Get RadianceHeader 313decodeHeader = do 314 sig <- getByteString $ B.length radianceFileSignature 315 when (sig /= radianceFileSignature) 316 (fail "Invalid radiance file signature") 317 318 infos <- decodeInfos 319 let formatKey = BC.pack "FORMAT" 320 case partition (\(k,_) -> k /= formatKey) infos of 321 (_, []) -> fail "No radiance format specified" 322 (info, [(_, formatString)]) -> 323 case runGet get $ L.fromChunks [formatString] of 324 Left err -> fail err 325 Right format -> do 326 (n1, n2, b) <- (,,) <$> decodeNum 327 <*> decodeNum 328 <*> getRemainingBytes 329 return . RadianceHeader info format n1 n2 $ L.fromChunks [b] 330 331 _ -> fail "Multiple radiance format specified" 332 333toFloat :: RGBE -> PixelRGBF 334toFloat (RGBE r g b e) = PixelRGBF rf gf bf 335 where f = encodeFloat 1 $ fromIntegral e - (128 + 8) 336 rf = (fromIntegral r + 0.0) * f 337 gf = (fromIntegral g + 0.0) * f 338 bf = (fromIntegral b + 0.0) * f 339 340encodeScanlineColor :: M.STVector s Word8 341 -> M.STVector s Word8 342 -> Int 343 -> ST s Int 344encodeScanlineColor vec outVec outIdx = do 345 val <- vec `M.unsafeRead` 0 346 runLength 1 0 val 1 outIdx 347 where maxIndex = M.length vec 348 349 pushRun len val at = do 350 (outVec `M.unsafeWrite` at) $ fromIntegral $ len .|. 0x80 351 (outVec `M.unsafeWrite` (at + 1)) val 352 return $ at + 2 353 354 pushData start len at = do 355 (outVec `M.unsafeWrite` at) $ fromIntegral len 356 let first = start - len 357 end = start - 1 358 offset = at - first + 1 359 forM_ [first .. end] $ \i -> do 360 v <- vec `M.unsafeRead` i 361 (outVec `M.unsafeWrite` (offset + i)) v 362 363 return $ at + len + 1 364 365 -- End of scanline, empty the thing 366 runLength run cpy prev idx at | idx >= maxIndex = 367 case (run, cpy) of 368 (0, 0) -> pure at 369 (0, n) -> pushData idx n at 370 (n, 0) -> pushRun n prev at 371 (_, _) -> error "HDR - Run length algorithm is wrong" 372 373 -- full runlength, we must write the packet 374 runLength r@127 _ prev idx at = do 375 val <- vec `M.unsafeRead` idx 376 pushRun r prev at >>= 377 runLength 1 0 val (idx + 1) 378 379 -- full copy, we must write the packet 380 runLength _ c@127 _ idx at = do 381 val <- vec `M.unsafeRead` idx 382 pushData idx c at >>= 383 runLength 1 0 val (idx + 1) 384 385 runLength n 0 prev idx at = do 386 val <- vec `M.unsafeRead` idx 387 case val == prev of 388 True -> runLength (n + 1) 0 prev (idx + 1) at 389 False | n < 4 -> runLength 0 (n + 1) val (idx + 1) at 390 False -> 391 pushRun n prev at >>= 392 runLength 1 0 val (idx + 1) 393 394 runLength 0 n prev idx at = do 395 val <- vec `M.unsafeRead` idx 396 if val /= prev 397 then runLength 0 (n + 1) val (idx + 1) at 398 else 399 pushData (idx - 1) (n - 1) at >>= 400 runLength (2 :: Int) 0 val (idx + 1) 401 402 runLength _ _ _ _ _ = 403 error "HDR RLE inconsistent state" 404 405-- | Write an High dynamic range image into a radiance 406-- image file on disk. 407writeHDR :: FilePath -> Image PixelRGBF -> IO () 408writeHDR filename img = L.writeFile filename $ encodeHDR img 409 410-- | Write a RLE encoded High dynamic range image into a radiance 411-- image file on disk. 412writeRLENewStyleHDR :: FilePath -> Image PixelRGBF -> IO () 413writeRLENewStyleHDR filename img = 414 L.writeFile filename $ encodeRLENewStyleHDR img 415 416-- | Encode an High dynamic range image into a radiance image 417-- file format. 418-- Alias for encodeRawHDR 419encodeHDR :: Image PixelRGBF -> L.ByteString 420encodeHDR = encodeRawHDR 421 422-- | Encode an High dynamic range image into a radiance image 423-- file format. without compression 424encodeRawHDR :: Image PixelRGBF -> L.ByteString 425encodeRawHDR pic = encode descriptor 426 where 427 newImage = pixelMap rgbeInRgba pic 428 -- we are cheating to death here, the layout we want 429 -- correspond to the layout of pixelRGBA8, so we 430 -- convert 431 rgbeInRgba pixel = PixelRGBA8 r g b e 432 where RGBE r g b e = toRGBE pixel 433 434 descriptor = RadianceHeader 435 { radianceInfos = [] 436 , radianceFormat = FormatRGBE 437 , radianceHeight = imageHeight pic 438 , radianceWidth = imageWidth pic 439 , radianceData = L.fromChunks [toByteString $ imageData newImage] 440 } 441 442 443-- | Encode an High dynamic range image into a radiance image 444-- file format using a light RLE compression. Some problems 445-- seem to arise with some image viewer. 446encodeRLENewStyleHDR :: Image PixelRGBF -> L.ByteString 447encodeRLENewStyleHDR pic = encode $ runST $ do 448 let w = imageWidth pic 449 h = imageHeight pic 450 451 scanLineR <- M.new w :: ST s (M.STVector s Word8) 452 scanLineG <- M.new w 453 scanLineB <- M.new w 454 scanLineE <- M.new w 455 456 encoded <- 457 forM [0 .. h - 1] $ \line -> do 458 buff <- M.new $ w * 4 + w `div` 127 + 2 459 let columner col | col >= w = return () 460 columner col = do 461 let RGBE r g b e = toRGBE $ pixelAt pic col line 462 (scanLineR `M.unsafeWrite` col) r 463 (scanLineG `M.unsafeWrite` col) g 464 (scanLineB `M.unsafeWrite` col) b 465 (scanLineE `M.unsafeWrite` col) e 466 467 columner (col + 1) 468 469 columner 0 470 471 (buff `M.unsafeWrite` 0) 2 472 (buff `M.unsafeWrite` 1) 2 473 (buff `M.unsafeWrite` 2) $ fromIntegral ((w .>>. 8) .&. 0xFF) 474 (buff `M.unsafeWrite` 3) $ fromIntegral (w .&. 0xFF) 475 476 i1 <- encodeScanlineColor scanLineR buff 4 477 i2 <- encodeScanlineColor scanLineG buff i1 478 i3 <- encodeScanlineColor scanLineB buff i2 479 endIndex <- encodeScanlineColor scanLineE buff i3 480 481 (\v -> blitVector v 0 endIndex) <$> V.unsafeFreeze buff 482 483 pure RadianceHeader 484 { radianceInfos = [] 485 , radianceFormat = FormatRGBE 486 , radianceHeight = h 487 , radianceWidth = w 488 , radianceData = L.fromChunks encoded 489 } 490 491 492decodeRadiancePicture :: RadianceHeader -> HDRReader s (MutableImage s PixelRGBF) 493decodeRadiancePicture hdr = do 494 let width = abs $ radianceWidth hdr 495 height = abs $ radianceHeight hdr 496 packedData = radianceData hdr 497 498 scanLine <- lift $ M.new $ width * 4 499 resultBuffer <- lift $ M.new $ width * height * 3 500 501 let scanLineImage = MutableImage 502 { mutableImageWidth = width 503 , mutableImageHeight = 1 504 , mutableImageData = scanLine 505 } 506 507 finalImage = MutableImage 508 { mutableImageWidth = width 509 , mutableImageHeight = height 510 , mutableImageData = resultBuffer 511 } 512 513 let scanLineExtractor readIdx line = do 514 let color = unpackColor packedData readIdx 515 inner | isNewRunLengthMarker color = do 516 let calcSize = checkLineLength color 517 when (calcSize /= width) 518 (throwE "Invalid sanline size") 519 pure $ \idx -> newStyleRLE packedData (idx + 4) 520 | otherwise = pure $ oldStyleRLE packedData 521 f <- inner 522 newRead <- f readIdx scanLine 523 forM_ [0 .. width - 1] $ \i -> do 524 -- mokay, it's a hack, but I don't want to define a 525 -- pixel instance of RGBE... 526 PixelRGBA8 r g b e <- lift $ readPixel scanLineImage i 0 527 lift $ writePixel finalImage i line . toFloat $ RGBE r g b e 528 529 return newRead 530 531 foldM_ scanLineExtractor 0 [0 .. height - 1] 532 533 return finalImage 534 535