diff --git a/Codec/Picture/BitWriter.hs b/Codec/Picture/BitWriter.hs --- a/Codec/Picture/BitWriter.hs +++ b/Codec/Picture/BitWriter.hs @@ -2,7 +2,8 @@ -- | This module implement helper functions to read & write data -- at bits level. module Codec.Picture.BitWriter( BoolReader - , BoolState( .. ) + , emptyBoolState + , BoolState , byteAlignJpg , getNextBitsLSBFirst , getNextBitsMSBFirst @@ -15,6 +16,11 @@ , newWriteStateRef , finalizeBoolWriter , writeBits' + + , initBoolState + , initBoolStateJpg + , execBoolReader + , runBoolReaderWith ) where import Data.STRef @@ -39,12 +45,36 @@ {-# UNPACK #-} !Word8 !B.ByteString +emptyBoolState :: BoolState +emptyBoolState = BoolState (-1) 0 B.empty + -- | Type used to read bits type BoolReader s a = S.StateT BoolState (ST s) a runBoolReader :: BoolReader s a -> ST s a runBoolReader action = S.evalStateT action $ BoolState 0 0 B.empty +runBoolReaderWith :: BoolState -> BoolReader s a -> ST s (a, BoolState) +runBoolReaderWith st action = S.runStateT action st + +execBoolReader :: BoolState -> BoolReader s a -> ST s BoolState +execBoolReader st reader = S.execStateT reader st + +initBoolState :: B.ByteString -> BoolState +initBoolState str = case B.uncons str of + Nothing -> BoolState 0 0 B.empty + Just (v, rest) -> BoolState 0 v rest + +initBoolStateJpg :: B.ByteString -> BoolState +initBoolStateJpg str = + case B.uncons str of + Nothing -> BoolState 0 0 B.empty + Just (0xFF, rest) -> case B.uncons rest of + Nothing -> BoolState maxBound 0 B.empty + Just (0x00, afterMarker) -> BoolState 7 0xFF afterMarker + Just (_ , afterMarker) -> initBoolStateJpg afterMarker + Just (v, rest) -> BoolState 7 v rest + -- | Bitify a list of things to decode. setDecodedString :: B.ByteString -> BoolReader s () setDecodedString str = case B.uncons str of @@ -115,9 +145,11 @@ Nothing -> S.put $ BoolState maxBound 0 B.empty Just (0xFF, rest) -> case B.uncons rest of Nothing -> S.put $ BoolState maxBound 0 B.empty - Just (0x00, afterMarker) -> S.put $ BoolState 7 0xFF afterMarker + Just (0x00, afterMarker) -> -- trace "00" $ + S.put $ BoolState 7 0xFF afterMarker Just (_ , afterMarker) -> setDecodedStringJpg afterMarker - Just (v, rest) -> S.put $ BoolState 7 v rest + Just (v, rest) -> + S.put $ BoolState 7 v rest -------------------------------------------------- ---- Writer diff --git a/Codec/Picture/Bitmap.hs b/Codec/Picture/Bitmap.hs --- a/Codec/Picture/Bitmap.hs +++ b/Codec/Picture/Bitmap.hs @@ -27,6 +27,8 @@ import Data.Binary.Get( Get , getWord16le , getWord32le + , bytesRead + , skip ) import Data.Word( Word32, Word16, Word8 ) @@ -257,15 +259,21 @@ -- decodeBitmap :: B.ByteString -> Either String DynamicImage decodeBitmap str = flip runGetStrict str $ do - _hdr <- get :: Get BmpHeader + hdr <- get :: Get BmpHeader bmpHeader <- get :: Get BmpInfoHeader case (bitPerPixel bmpHeader, planes bmpHeader, bitmapCompression bmpHeader) of -- (32, 1, 0) -> {- ImageRGBA8 <$>-} fail "Meuh" (24, 1, 0) -> do + readed <- bytesRead + + when (readed > fromIntegral (dataOffset hdr)) + (fail "Invalid bmp image, data in header") + + skip . fromIntegral $ dataOffset hdr - fromIntegral readed rest <- getRemainingBytes return . ImageRGB8 $ decodeImageRGB8 bmpHeader rest - _ -> fail "Can't handle BMP file" + a -> fail $ "Can't handle BMP file " ++ show a -- | Write an image in a file use the bitmap format. diff --git a/Codec/Picture/Jpg.hs b/Codec/Picture/Jpg.hs --- a/Codec/Picture/Jpg.hs +++ b/Codec/Picture/Jpg.hs @@ -5,1303 +5,716 @@ {-# LANGUAGE BangPatterns #-} {-# OPTIONS_GHC -fspec-constr-count=5 #-} -- | Module used for JPEG file loading and writing. -module Codec.Picture.Jpg( decodeJpeg, encodeJpegAtQuality, encodeJpeg ) where - -import Control.Arrow( (>>>) ) -import Control.Applicative( (<$>), (<*>)) -import Control.Monad( when, replicateM, forM, forM_, foldM_, unless ) -import Control.Monad.ST( ST, runST ) -import Control.Monad.Trans( lift ) -import qualified Control.Monad.Trans.State.Strict as S - -import Data.Maybe( fromMaybe ) -import Data.List( find, foldl' ) -import Data.Bits( (.|.), (.&.), unsafeShiftL, unsafeShiftR ) -import Data.Int( Int16, Int32 ) -import Data.Word(Word8, Word16, Word32) -import Data.Binary( Binary(..), encode ) - -import Data.Binary.Get( Get - , getWord8 - , getWord16be - , getByteString - , skip - , bytesRead - ) - -import Data.Binary.Put( Put - , putWord8 - , putWord16be - , putLazyByteString - ) - -import Data.Maybe( fromJust ) -import qualified Data.Vector as V -import Data.Vector.Unboxed( (!) ) -import qualified Data.Vector.Unboxed as VU -import qualified Data.Vector.Storable as VS -import qualified Data.Vector.Storable.Mutable as M --- import Data.Array.Unboxed( Array, UArray, elems, listArray, (!) ) -import qualified Data.ByteString as B -import qualified Data.ByteString.Lazy as L -import Foreign.Storable ( Storable ) - -import Codec.Picture.InternalHelper -import Codec.Picture.BitWriter -import Codec.Picture.Types -import Codec.Picture.Jpg.Types -import Codec.Picture.Jpg.DefaultTable -import Codec.Picture.Jpg.FastIdct -import Codec.Picture.Jpg.FastDct - --------------------------------------------------- ----- Types --------------------------------------------------- -data JpgFrameKind = - JpgBaselineDCTHuffman - | JpgExtendedSequentialDCTHuffman - | JpgProgressiveDCTHuffman - | JpgLosslessHuffman - | JpgDifferentialSequentialDCTHuffman - | JpgDifferentialProgressiveDCTHuffman - | JpgDifferentialLosslessHuffman - | JpgExtendedSequentialArithmetic - | JpgProgressiveDCTArithmetic - | JpgLosslessArithmetic - | JpgDifferentialSequentialDCTArithmetic - | JpgDifferentialProgressiveDCTArithmetic - | JpgDifferentialLosslessArithmetic - | JpgQuantizationTable - | JpgHuffmanTableMarker - | JpgStartOfScan - | JpgEndOfImage - | JpgAppSegment Word8 - | JpgExtensionSegment Word8 - - | JpgRestartInterval - | JpgRestartIntervalEnd Word8 - deriving (Eq, Show) - -type HuffmanTreeInfo = HuffmanPackedTree - -data JpgFrame = - JpgAppFrame !Word8 B.ByteString - | JpgExtension !Word8 B.ByteString - | JpgQuantTable ![JpgQuantTableSpec] - | JpgHuffmanTable ![(JpgHuffmanTableSpec, HuffmanTreeInfo)] - | JpgScanBlob !JpgScanHeader !L.ByteString - | JpgScans !JpgFrameKind !JpgFrameHeader - | JpgIntervalRestart !Word16 - deriving Show - -data JpgFrameHeader = JpgFrameHeader - { jpgFrameHeaderLength :: !Word16 - , jpgSamplePrecision :: !Word8 - , jpgHeight :: !Word16 - , jpgWidth :: !Word16 - , jpgImageComponentCount :: !Word8 - , jpgComponents :: ![JpgComponent] - } - deriving Show - -instance SizeCalculable JpgFrameHeader where - calculateSize hdr = 2 + 1 + 2 + 2 + 1 - + sum [calculateSize c | c <- jpgComponents hdr] - -data JpgComponent = JpgComponent - { componentIdentifier :: !Word8 - -- | Stored with 4 bits - , horizontalSamplingFactor :: !Word8 - -- | Stored with 4 bits - , verticalSamplingFactor :: !Word8 - , quantizationTableDest :: !Word8 - } - deriving Show - -instance SizeCalculable JpgComponent where - calculateSize _ = 3 - -data JpgImage = JpgImage { jpgFrame :: [JpgFrame]} - deriving Show - -data JpgScanSpecification = JpgScanSpecification - { componentSelector :: !Word8 - -- | Encoded as 4 bits - , dcEntropyCodingTable :: !Word8 - -- | Encoded as 4 bits - , acEntropyCodingTable :: !Word8 - - } - deriving Show - -instance SizeCalculable JpgScanSpecification where - calculateSize _ = 2 - -data JpgScanHeader = JpgScanHeader - { scanLength :: !Word16 - , scanComponentCount :: !Word8 - , scans :: [JpgScanSpecification] - - -- | (begin, end) - , spectralSelection :: (Word8, Word8) - - -- | Encoded as 4 bits - , successiveApproxHigh :: !Word8 - - -- | Encoded as 4 bits - , successiveApproxLow :: !Word8 - } - deriving Show - -instance SizeCalculable JpgScanHeader where - calculateSize hdr = 2 + 1 - + sum [calculateSize c | c <- scans hdr] - + 2 - + 1 - -data JpgQuantTableSpec = JpgQuantTableSpec - { -- | Stored on 4 bits - quantPrecision :: !Word8 - - -- | Stored on 4 bits - , quantDestination :: !Word8 - - , quantTable :: MacroBlock Int16 - } - deriving Show - --- | Type introduced only to avoid some typeclass overlapping --- problem -newtype TableList a = TableList [a] - -class SizeCalculable a where - calculateSize :: a -> Int - -instance (SizeCalculable a, Binary a) => Binary (TableList a) where - put (TableList lst) = do - putWord16be . fromIntegral $ sum [calculateSize table | table <- lst] + 2 - mapM_ put lst - - get = TableList <$> (getWord16be >>= \s -> innerParse (fromIntegral s - 2)) - where innerParse :: Int -> Get [a] - innerParse 0 = return [] - innerParse size = do - onStart <- fromIntegral <$> bytesRead - table <- get - onEnd <- fromIntegral <$> bytesRead - (table :) <$> innerParse (size - (onEnd - onStart)) - -instance SizeCalculable JpgQuantTableSpec where - calculateSize table = - 1 + (fromIntegral (quantPrecision table) + 1) * 64 - -instance Binary JpgQuantTableSpec where - put table = do - let precision = quantPrecision table - put4BitsOfEach precision (quantDestination table) - forM_ (VS.toList $ quantTable table) $ \coeff -> - if precision == 0 then putWord8 $ fromIntegral coeff - else putWord16be $ fromIntegral coeff - - get = do - (precision, dest) <- get4BitOfEach - coeffs <- replicateM 64 $ if precision == 0 - then fromIntegral <$> getWord8 - else fromIntegral <$> getWord16be - return JpgQuantTableSpec - { quantPrecision = precision - , quantDestination = dest - , quantTable = VS.fromListN 64 coeffs - } - -data JpgHuffmanTableSpec = JpgHuffmanTableSpec - { -- | 0 : DC, 1 : AC, stored on 4 bits - huffmanTableClass :: !DctComponent - -- | Stored on 4 bits - , huffmanTableDest :: !Word8 - - , huffSizes :: !(VU.Vector Word8) - , huffCodes :: !(V.Vector (VU.Vector Word8)) - } - deriving Show - -buildPackedHuffmanTree :: V.Vector (VU.Vector Word8) -> HuffmanTree -buildPackedHuffmanTree = buildHuffmanTree . map VU.toList . V.toList - -huffmanPackedDecode :: HuffmanPackedTree -> BoolReader s Word8 -huffmanPackedDecode table = getNextBitJpg >>= aux 0 - where aux idx b | (v .&. 0x8000) /= 0 = return 0 - | (v .&. 0x4000) /= 0 = return . fromIntegral $ v .&. 0xFF - | otherwise = getNextBitJpg >>= aux v - where tableIndex | b = idx + 1 - | otherwise = idx - v = table `VS.unsafeIndex` fromIntegral tableIndex - --------------------------------------------------- ----- Serialization instances --------------------------------------------------- -commonMarkerFirstByte :: Word8 -commonMarkerFirstByte = 0xFF - -checkMarker :: Word8 -> Word8 -> Get () -checkMarker b1 b2 = do - rb1 <- getWord8 - rb2 <- getWord8 - when (rb1 /= b1 || rb2 /= b2) - (fail "Invalid marker used") - -eatUntilCode :: Get () -eatUntilCode = do - code <- getWord8 - unless (code == 0xFF) eatUntilCode - -instance SizeCalculable JpgHuffmanTableSpec where - calculateSize table = 1 + 16 + sum [fromIntegral e | e <- VU.toList $ huffSizes table] - -instance Binary JpgHuffmanTableSpec where - put table = do - let classVal = if huffmanTableClass table == DcComponent - then 0 else 1 - put4BitsOfEach classVal $ huffmanTableDest table - mapM_ put . VU.toList $ huffSizes table - forM_ [0 .. 15] $ \i -> - when (huffSizes table ! i /= 0) - (let elements = VU.toList $ huffCodes table V.! i - in mapM_ put elements) - - get = do - (huffClass, huffDest) <- get4BitOfEach - sizes <- replicateM 16 getWord8 - codes <- forM sizes $ \s -> - VU.replicateM (fromIntegral s) getWord8 - return JpgHuffmanTableSpec - { huffmanTableClass = - if huffClass == 0 then DcComponent else AcComponent - , huffmanTableDest = huffDest - , huffSizes = VU.fromListN 16 sizes - , huffCodes = V.fromListN 16 codes - } - -instance Binary JpgImage where - put (JpgImage { jpgFrame = frames }) = - putWord8 0xFF >> putWord8 0xD8 >> mapM_ putFrame frames - >> putWord8 0xFF >> putWord8 0xD9 - - get = do - let startOfImageMarker = 0xD8 - -- endOfImageMarker = 0xD9 - checkMarker commonMarkerFirstByte startOfImageMarker - eatUntilCode - frames <- parseFrames - {-checkMarker commonMarkerFirstByte endOfImageMarker-} - return JpgImage { jpgFrame = frames } - -takeCurrentFrame :: Get B.ByteString -takeCurrentFrame = do - size <- getWord16be - getByteString (fromIntegral size - 2) - -putFrame :: JpgFrame -> Put -putFrame (JpgAppFrame appCode str) = - put (JpgAppSegment appCode) >> putWord16be (fromIntegral $ B.length str) >> put str -putFrame (JpgExtension appCode str) = - put (JpgExtensionSegment appCode) >> putWord16be (fromIntegral $ B.length str) >> put str -putFrame (JpgQuantTable tables) = - put JpgQuantizationTable >> put (TableList tables) -putFrame (JpgHuffmanTable tables) = - put JpgHuffmanTableMarker >> put (TableList $ map fst tables) -putFrame (JpgIntervalRestart size) = - put JpgRestartInterval >> put (RestartInterval size) -putFrame (JpgScanBlob hdr blob) = - put JpgStartOfScan >> put hdr >> putLazyByteString blob -putFrame (JpgScans kind hdr) = - put kind >> put hdr - -extractScanContent :: L.ByteString -> (L.ByteString, L.ByteString) -extractScanContent str = aux 0 - where maxi = fromIntegral $ L.length str - 1 - - aux n | n >= maxi = (str, L.empty) - | v == 0xFF && vNext /= 0 && not isReset = L.splitAt n str - | otherwise = aux (n + 1) - where v = str `L.index` n - vNext = str `L.index` (n + 1) - isReset = 0xD0 <= vNext && vNext <= 0xD7 - -parseFrames :: Get [JpgFrame] -parseFrames = do - kind <- get - let parseNextFrame = do - word <- getWord8 - when (word /= 0xFF) $ do - readedData <- bytesRead - fail $ "Invalid Frame marker (" ++ show word - ++ ", bytes read : " ++ show readedData ++ ")" - parseFrames - - case kind of - JpgEndOfImage -> return [] - JpgAppSegment c -> - (\frm lst -> JpgAppFrame c frm : lst) <$> takeCurrentFrame <*> parseNextFrame - JpgExtensionSegment c -> - (\frm lst -> JpgExtension c frm : lst) <$> takeCurrentFrame <*> parseNextFrame - JpgQuantizationTable -> - (\(TableList quants) lst -> JpgQuantTable quants : lst) <$> get <*> parseNextFrame - JpgRestartInterval -> - (\(RestartInterval i) lst -> JpgIntervalRestart i : lst) <$> get <*> parseNextFrame - JpgHuffmanTableMarker -> - (\(TableList huffTables) lst -> - JpgHuffmanTable [(t, packHuffmanTree . buildPackedHuffmanTree $ huffCodes t) | t <- huffTables] : lst) - <$> get <*> parseNextFrame - JpgStartOfScan -> - (\frm imgData -> - let (d, other) = extractScanContent imgData - in - case runGet parseFrames (L.drop 1 other) of - Left _ -> [JpgScanBlob frm d] - Right lst -> JpgScanBlob frm d : lst - ) <$> get <*> getRemainingLazyBytes - - _ -> (\hdr lst -> JpgScans kind hdr : lst) <$> get <*> parseNextFrame - -secondStartOfFrameByteOfKind :: JpgFrameKind -> Word8 -secondStartOfFrameByteOfKind = aux - where - aux JpgBaselineDCTHuffman = 0xC0 - aux JpgExtendedSequentialDCTHuffman = 0xC1 - aux JpgProgressiveDCTHuffman = 0xC2 - aux JpgLosslessHuffman = 0xC3 - aux JpgDifferentialSequentialDCTHuffman = 0xC5 - aux JpgDifferentialProgressiveDCTHuffman = 0xC6 - aux JpgDifferentialLosslessHuffman = 0xC7 - aux JpgExtendedSequentialArithmetic = 0xC9 - aux JpgProgressiveDCTArithmetic = 0xCA - aux JpgLosslessArithmetic = 0xCB - aux JpgHuffmanTableMarker = 0xC4 - aux JpgDifferentialSequentialDCTArithmetic = 0xCD - aux JpgDifferentialProgressiveDCTArithmetic = 0xCE - aux JpgDifferentialLosslessArithmetic = 0xCF - aux JpgEndOfImage = 0xD9 - aux JpgQuantizationTable = 0xDB - aux JpgStartOfScan = 0xDA - aux JpgRestartInterval = 0xDD - aux (JpgRestartIntervalEnd v) = v - aux (JpgAppSegment a) = a - aux (JpgExtensionSegment a) = a - -data JpgImageKind = BaseLineDCT | ProgressiveDCT - -instance Binary JpgFrameKind where - put v = putWord8 0xFF >> put (secondStartOfFrameByteOfKind v) - get = do - -- no lookahead :( - {-word <- getWord8-} - word2 <- getWord8 - return $ case word2 of - 0xC0 -> JpgBaselineDCTHuffman - 0xC1 -> JpgExtendedSequentialDCTHuffman - 0xC2 -> JpgProgressiveDCTHuffman - 0xC3 -> JpgLosslessHuffman - 0xC4 -> JpgHuffmanTableMarker - 0xC5 -> JpgDifferentialSequentialDCTHuffman - 0xC6 -> JpgDifferentialProgressiveDCTHuffman - 0xC7 -> JpgDifferentialLosslessHuffman - 0xC9 -> JpgExtendedSequentialArithmetic - 0xCA -> JpgProgressiveDCTArithmetic - 0xCB -> JpgLosslessArithmetic - 0xCD -> JpgDifferentialSequentialDCTArithmetic - 0xCE -> JpgDifferentialProgressiveDCTArithmetic - 0xCF -> JpgDifferentialLosslessArithmetic - 0xD9 -> JpgEndOfImage - 0xDA -> JpgStartOfScan - 0xDB -> JpgQuantizationTable - 0xDD -> JpgRestartInterval - a | a >= 0xF0 -> JpgExtensionSegment a - | a >= 0xE0 -> JpgAppSegment a - | a >= 0xD0 && a <= 0xD7 -> JpgRestartIntervalEnd a - | otherwise -> error ("Invalid frame marker (" ++ show a ++ ")") - -put4BitsOfEach :: Word8 -> Word8 -> Put -put4BitsOfEach a b = put $ (a `unsafeShiftL` 4) .|. b - -get4BitOfEach :: Get (Word8, Word8) -get4BitOfEach = do - val <- get - return ((val `unsafeShiftR` 4) .&. 0xF, val .&. 0xF) - -newtype RestartInterval = RestartInterval Word16 - -instance Binary RestartInterval where - put (RestartInterval i) = putWord16be 4 >> putWord16be i - get = do - size <- getWord16be - when (size /= 4) (fail "Invalid jpeg restart interval size") - RestartInterval <$> getWord16be - -instance Binary JpgComponent where - get = do - ident <- getWord8 - (horiz, vert) <- get4BitOfEach - quantTableIndex <- getWord8 - return JpgComponent - { componentIdentifier = ident - , horizontalSamplingFactor = horiz - , verticalSamplingFactor = vert - , quantizationTableDest = quantTableIndex - } - put v = do - put $ componentIdentifier v - put4BitsOfEach (horizontalSamplingFactor v) $ verticalSamplingFactor v - put $ quantizationTableDest v - -instance Binary JpgFrameHeader where - get = do - beginOffset <- fromIntegral <$> bytesRead - frmHLength <- getWord16be - samplePrec <- getWord8 - h <- getWord16be - w <- getWord16be - compCount <- getWord8 - components <- replicateM (fromIntegral compCount) get - endOffset <- fromIntegral <$> bytesRead - when (beginOffset - endOffset < fromIntegral frmHLength) - (skip $ fromIntegral frmHLength - (endOffset - beginOffset)) - return JpgFrameHeader - { jpgFrameHeaderLength = frmHLength - , jpgSamplePrecision = samplePrec - , jpgHeight = h - , jpgWidth = w - , jpgImageComponentCount = compCount - , jpgComponents = components - } - - put v = do - putWord16be $ jpgFrameHeaderLength v - putWord8 $ jpgSamplePrecision v - putWord16be $ jpgHeight v - putWord16be $ jpgWidth v - putWord8 $ jpgImageComponentCount v - mapM_ put $ jpgComponents v - -instance Binary JpgScanSpecification where - put v = do - put $ componentSelector v - put4BitsOfEach (dcEntropyCodingTable v) $ acEntropyCodingTable v - - get = do - compSel <- get - (dc, ac) <- get4BitOfEach - return JpgScanSpecification { - componentSelector = compSel - , dcEntropyCodingTable = dc - , acEntropyCodingTable = ac - } - -instance Binary JpgScanHeader where - get = do - thisScanLength <- getWord16be - compCount <- getWord8 - comp <- replicateM (fromIntegral compCount) get - specBeg <- get - specEnd <- get - (approxHigh, approxLow) <- get4BitOfEach - - return JpgScanHeader { - scanLength = thisScanLength, - scanComponentCount = compCount, - scans = comp, - spectralSelection = (specBeg, specEnd), - successiveApproxHigh = approxHigh, - successiveApproxLow = approxLow - } - - put v = do - putWord16be $ scanLength v - putWord8 $ scanComponentCount v - mapM_ put $ scans v - putWord8 . fst $ spectralSelection v - putWord8 . snd $ spectralSelection v - put4BitsOfEach (successiveApproxHigh v) $ successiveApproxLow v - -quantize :: MacroBlock Int16 -> MutableMacroBlock s Int32 - -> ST s (MutableMacroBlock s Int32) -quantize table block = update 0 - where update 64 = return block - update idx = do - val <- block `M.unsafeRead` idx - let q = fromIntegral (table `VS.unsafeIndex` idx) - finalValue = (val + (q `div` 2)) `quot` q -- rounded integer division - (block `M.unsafeWrite` idx) finalValue - update $ idx + 1 - --- | Apply a quantization matrix to a macroblock -{-# INLINE deQuantize #-} -deQuantize :: MacroBlock Int16 -> MutableMacroBlock s Int16 - -> ST s (MutableMacroBlock s Int16) -deQuantize table block = update 0 - where update 64 = return block - update i = do - val <- block `M.unsafeRead` i - let finalValue = val * (table `VS.unsafeIndex` i) - (block `M.unsafeWrite` i) finalValue - update $ i + 1 - -inverseDirectCosineTransform :: MutableMacroBlock s Int16 - -> ST s (MutableMacroBlock s Int16) -inverseDirectCosineTransform mBlock = - fastIdct mBlock >>= mutableLevelShift - -zigZagOrder :: MacroBlock Int -zigZagOrder = makeMacroBlock $ concat - [[ 0, 1, 5, 6,14,15,27,28] - ,[ 2, 4, 7,13,16,26,29,42] - ,[ 3, 8,12,17,25,30,41,43] - ,[ 9,11,18,24,31,40,44,53] - ,[10,19,23,32,39,45,52,54] - ,[20,22,33,38,46,51,55,60] - ,[21,34,37,47,50,56,59,61] - ,[35,36,48,49,57,58,62,63] - ] - -zigZagReorderForwardv :: (Storable a, Num a) => VS.Vector a -> VS.Vector a -zigZagReorderForwardv vec = runST $ do - v <- M.new 64 - mv <- VS.thaw vec - zigZagReorderForward v mv >>= VS.freeze - -zigZagOrderForward :: MacroBlock Int -zigZagOrderForward = VS.generate 64 inv - where inv i = fromMaybe 0 $ VS.findIndex (i ==) zigZagOrder - -zigZagReorderForward :: (Storable a, Num a) - => MutableMacroBlock s a - -> MutableMacroBlock s a - -> ST s (MutableMacroBlock s a) -{-# SPECIALIZE INLINE zigZagReorderForward :: MutableMacroBlock s Int32 - -> MutableMacroBlock s Int32 - -> ST s (MutableMacroBlock s Int32) #-} -{-# SPECIALIZE INLINE zigZagReorderForward :: MutableMacroBlock s Int16 - -> MutableMacroBlock s Int16 - -> ST s (MutableMacroBlock s Int16) #-} -{-# SPECIALIZE INLINE zigZagReorderForward :: MutableMacroBlock s Word8 - -> MutableMacroBlock s Word8 - -> ST s (MutableMacroBlock s Word8) #-} -zigZagReorderForward zigzaged block = ordering zigZagOrderForward >> return zigzaged - where ordering !table = reorder (0 :: Int) - where reorder !i | i >= 64 = return () - reorder i = do - let idx = table `VS.unsafeIndex` i - v <- block `M.unsafeRead` idx - (zigzaged `M.unsafeWrite` i) v - reorder (i + 1) - -zigZagReorder :: MutableMacroBlock s Int16 -> MutableMacroBlock s Int16 - -> ST s (MutableMacroBlock s Int16) -zigZagReorder zigzaged block = do - let update i = do - let idx = zigZagOrder `VS.unsafeIndex` i - v <- block `M.unsafeRead` idx - (zigzaged `M.unsafeWrite` i) v - - reorder 63 = update 63 - reorder i = update i >> reorder (i + 1) - - reorder (0 :: Int) - return zigzaged - - --- | This is one of the most important function of the decoding, --- it form the barebone decoding pipeline for macroblock. It's all --- there is to know for macro block transformation -decodeMacroBlock :: MacroBlock DctCoefficients - -> MutableMacroBlock s Int16 - -> MutableMacroBlock s Int16 - -> ST s (MutableMacroBlock s Int16) -decodeMacroBlock quantizationTable zigZagBlock block = - deQuantize quantizationTable block >>= zigZagReorder zigZagBlock - >>= inverseDirectCosineTransform - -packInt :: [Bool] -> Int32 -packInt = foldl' bitStep 0 - where bitStep acc True = (acc `unsafeShiftL` 1) + 1 - bitStep acc False = acc `unsafeShiftL` 1 - --- | Unpack an int of the given size encoded from MSB to LSB. -unpackInt :: Int32 -> BoolReader s Int32 -unpackInt bitCount = packInt <$> replicateM (fromIntegral bitCount) getNextBitJpg - -powerOf :: Int32 -> Word32 -powerOf 0 = 0 -powerOf n = limit 1 0 - where val = abs n - limit range i | val < range = i - limit range i = limit (2 * range) (i + 1) - -encodeInt :: BoolWriteStateRef s -> Word32 -> Int32 -> ST s () -{-# INLINE encodeInt #-} -encodeInt st ssss n | n > 0 = writeBits' st (fromIntegral n) (fromIntegral ssss) -encodeInt st ssss n = writeBits' st (fromIntegral $ n - 1) (fromIntegral ssss) - -{-# INLINE decodeInt #-} -decodeInt :: Int32 -> BoolReader s Int32 -decodeInt ssss = do - signBit <- getNextBitJpg - let dataRange = 1 `unsafeShiftL` fromIntegral (ssss - 1) - leftBitCount = ssss - 1 - -- First following bits store the sign of the coefficient, and counted in - -- SSSS, so the bit count for the int, is ssss - 1 - if signBit - then (\w -> dataRange + fromIntegral w) <$> unpackInt leftBitCount - else (\w -> 1 - dataRange * 2 + fromIntegral w) <$> unpackInt leftBitCount - -dcCoefficientDecode :: HuffmanTreeInfo - -> BoolReader s DcCoefficient -dcCoefficientDecode dcTree = do - ssss <- huffmanPackedDecode dcTree - if ssss == 0 - then return 0 - else fromIntegral <$> decodeInt (fromIntegral ssss) - --- | Assume the macro block is initialized with zeroes -acCoefficientsDecode :: HuffmanTreeInfo -> MutableMacroBlock s Int16 - -> BoolReader s (MutableMacroBlock s Int16) -acCoefficientsDecode acTree mutableBlock = parseAcCoefficient 1 >> return mutableBlock - where parseAcCoefficient n | n >= 64 = return () - | otherwise = do - rrrrssss <- huffmanPackedDecode acTree - let rrrr = fromIntegral $ (rrrrssss `unsafeShiftR` 4) .&. 0xF - ssss = rrrrssss .&. 0xF - case (rrrr, ssss) of - ( 0, 0) -> return () - (0xF, 0) -> parseAcCoefficient (n + 16) - _ -> do - decoded <- fromIntegral <$> decodeInt (fromIntegral ssss) - lift $ (mutableBlock `M.unsafeWrite` (n + rrrr)) decoded - parseAcCoefficient (n + rrrr + 1) - -{- -decodeEOB :: HuffmanTreeInfo -> BoolReader s Int32 -decodeEOB tree = do - rrrrssss <- huffmanPackedDecode tree - -} - --- | Decompress a macroblock from a bitstream given the current configuration --- from the frame. -decompressMacroBlock :: HuffmanTreeInfo -- ^ Tree used for DC coefficient - -> HuffmanTreeInfo -- ^ Tree used for Ac coefficient - -> MacroBlock Int16 -- ^ Current quantization table - -> MutableMacroBlock s Int16 -- ^ A zigzag table, to avoid allocation - -> DcCoefficient -- ^ Previous dc value - -> BoolReader s (DcCoefficient, MutableMacroBlock s Int16) -decompressMacroBlock dcTree acTree quantizationTable zigzagBlock previousDc = do - dcDeltaCoefficient <- dcCoefficientDecode dcTree - block <- lift createEmptyMutableMacroBlock - let neoDcCoefficient = previousDc + dcDeltaCoefficient - lift $ (block `M.unsafeWrite` 0) neoDcCoefficient - fullBlock <- acCoefficientsDecode acTree block - decodedBlock <- lift $ decodeMacroBlock quantizationTable zigzagBlock fullBlock - return (neoDcCoefficient, decodedBlock) - -gatherQuantTables :: JpgImage -> [JpgQuantTableSpec] -gatherQuantTables img = concat [t | JpgQuantTable t <- jpgFrame img] - -gatherHuffmanTables :: JpgImage -> [(JpgHuffmanTableSpec, HuffmanTreeInfo)] -gatherHuffmanTables img = - concat [lst | JpgHuffmanTable lst <- jpgFrame img] ++ defaultHuffmanTables - - -gatherScanInfo :: JpgImage -> (JpgFrameKind, JpgFrameHeader) -gatherScanInfo img = fromJust $ unScan <$> find scanDesc (jpgFrame img) - where scanDesc (JpgScans _ _) = True - scanDesc _ = False - - unScan (JpgScans a b) = (a,b) - unScan _ = error "If this can happen, the JPEG image is ill-formed" - -pixelClamp :: Int16 -> Word8 -pixelClamp n = fromIntegral . min 255 $ max 0 n - -unpack444Y :: Int -- ^ x - -> Int -- ^ y - -> MutableImage s PixelYCbCr8 - -> MutableMacroBlock s Int16 - -> ST s () -unpack444Y x y (MutableImage { mutableImageWidth = imgWidth, mutableImageData = img }) - block = blockVert baseIdx 0 zero - where zero = 0 :: Int - baseIdx = x * 8 + y * 8 * imgWidth - - blockVert _ _ j | j >= 8 = return () - blockVert writeIdx readingIdx j = blockHoriz writeIdx readingIdx zero - where blockHoriz _ readIdx i | i >= 8 = blockVert (writeIdx + imgWidth) readIdx $ j + 1 - blockHoriz idx readIdx i = do - val <- pixelClamp <$> (block `M.unsafeRead` readIdx) - (img `M.unsafeWrite` idx) val - blockHoriz (idx + 1) (readIdx + 1) $ i + 1 - -unpack444Ycbcr :: Int -- ^ Component index - -> Int -- ^ x - -> Int -- ^ y - -> MutableImage s PixelYCbCr8 - -> MutableMacroBlock s Int16 - -> ST s () -unpack444Ycbcr compIdx x y - (MutableImage { mutableImageWidth = imgWidth, mutableImageData = img }) - block = blockVert baseIdx 0 zero - where zero = 0 :: Int - baseIdx = (x * 8 + y * 8 * imgWidth) * 3 + compIdx - - blockVert _ _ j | j >= 8 = return () - blockVert idx readIdx j = do - val0 <- pixelClamp <$> (block `M.unsafeRead` readIdx) - val1 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 1)) - val2 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 2)) - val3 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 3)) - val4 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 4)) - val5 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 5)) - val6 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 6)) - val7 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 7)) - - (img `M.unsafeWrite` idx) val0 - (img `M.unsafeWrite` (idx + (3 * 1))) val1 - (img `M.unsafeWrite` (idx + (3 * 2))) val2 - (img `M.unsafeWrite` (idx + (3 * 3))) val3 - (img `M.unsafeWrite` (idx + (3 * 4))) val4 - (img `M.unsafeWrite` (idx + (3 * 5))) val5 - (img `M.unsafeWrite` (idx + (3 * 6))) val6 - (img `M.unsafeWrite` (idx + (3 * 7))) val7 - - blockVert (idx + 3 * imgWidth) (readIdx + 8) $ j + 1 - - - {-where blockHoriz _ readIdx i | i >= 8 = blockVert (writeIdx + imgWidth * 3) readIdx $ j + 1-} - {-blockHoriz idx readIdx i = do-} - {-val <- pixelClamp <$> (block `M.unsafeRead` readIdx) -} - {-(img `M.unsafeWrite` idx) val-} - {-blockHoriz (idx + 3) (readIdx + 1) $ i + 1-} - -unpack422Ycbcr :: Int -- ^ Component index - -> Int -- ^ x - -> Int -- ^ y - -> MutableImage s PixelYCbCr8 - -> MutableMacroBlock s Int16 - -> ST s () -unpack422Ycbcr compIdx x y - (MutableImage { mutableImageWidth = imgWidth, - mutableImageHeight = _, mutableImageData = img }) - block = blockVert baseIdx 0 zero - where zero = 0 :: Int - baseIdx = (x * 8 * 2 + y * 8 * imgWidth) * 3 + compIdx - lineOffset = imgWidth * 3 - - blockVert _ _ j | j >= 8 = return () - blockVert idx readIdx j = do - v0 <- pixelClamp <$> (block `M.unsafeRead` readIdx) - v1 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 1)) - v2 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 2)) - v3 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 3)) - v4 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 4)) - v5 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 5)) - v6 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 6)) - v7 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 7)) - - (img `M.unsafeWrite` idx) v0 - (img `M.unsafeWrite` (idx + 3)) v0 - - (img `M.unsafeWrite` (idx + 6 * 1)) v1 - (img `M.unsafeWrite` (idx + 6 * 1 + 3)) v1 - - (img `M.unsafeWrite` (idx + 6 * 2)) v2 - (img `M.unsafeWrite` (idx + 6 * 2 + 3)) v2 - - (img `M.unsafeWrite` (idx + 6 * 3)) v3 - (img `M.unsafeWrite` (idx + 6 * 3 + 3)) v3 - - (img `M.unsafeWrite` (idx + 6 * 4)) v4 - (img `M.unsafeWrite` (idx + 6 * 4 + 3)) v4 - - (img `M.unsafeWrite` (idx + 6 * 5)) v5 - (img `M.unsafeWrite` (idx + 6 * 5 + 3)) v5 - - (img `M.unsafeWrite` (idx + 6 * 6)) v6 - (img `M.unsafeWrite` (idx + 6 * 6 + 3)) v6 - - (img `M.unsafeWrite` (idx + 6 * 7)) v7 - (img `M.unsafeWrite` (idx + 6 * 7 + 3)) v7 - - blockVert (idx + lineOffset) (readIdx + 8) $ j + 1 - --- | Given a size coefficient (how much a pixel span horizontally --- and vertically), the position of the macroblock, return a list --- of indices and value to be stored in an array (like the final --- image) -unpackMacroBlock :: Int -- ^ Component count - -> Int -- ^ Component index - -> Int -- ^ Width coefficient - -> Int -- ^ Height coefficient - -> Int -- ^ x - -> Int -- ^ y - -> MutableImage s PixelYCbCr8 - -> MutableMacroBlock s Int16 - -> ST s () -unpackMacroBlock compCount compIdx wCoeff hCoeff x y - (MutableImage { mutableImageWidth = imgWidth, - mutableImageHeight = imgHeight, mutableImageData = img }) - block = -- trace (printf "w:%d h:%d x:%d y:%d wCoeff:%d hCoeff:%d" imgWidth imgHeight x y wCoeff hCoeff) $ - blockVert 0 - where blockVert j | j >= 8 = return () - blockVert j = blockHoriz 0 - where yBase = (y * 8 + j) * hCoeff - blockHoriz i | i >= 8 = blockVert $ j + 1 - blockHoriz i = (pixelClamp <$> (block `M.unsafeRead` (i + j * 8))) >>= horizDup 0 - where xBase = (x * 8 + i) * wCoeff - horizDup wDup _ | wDup >= wCoeff = blockHoriz $ i + 1 - horizDup wDup compVal = vertDup 0 - where vertDup hDup | hDup >= hCoeff = horizDup (wDup + 1) compVal - vertDup hDup = do - let xPos = xBase + wDup - yPos = yBase + hDup - - when (xPos < imgWidth && yPos < imgHeight) - (do let mutableIdx = (xPos + yPos * imgWidth) * compCount + compIdx - (img `M.unsafeWrite` mutableIdx) compVal) - - vertDup $ hDup + 1 - --- | Type only used to make clear what kind of integer we are carrying --- Might be transformed into newtype in the future -type DcCoefficient = Int16 - --- | Same as for DcCoefficient, to provide nicer type signatures -type DctCoefficients = DcCoefficient - -decodeRestartInterval :: BoolReader s Int32 -decodeRestartInterval = return (-1) {- do - bits <- replicateM 8 getNextBitJpg - if bits == replicate 8 True - then do - marker <- replicateM 8 getNextBitJpg - return $ packInt marker - else return (-1) - -} - -decodeImage :: JpgImage - -> Int -- ^ Component count - -> MutableImage s PixelYCbCr8 -- ^ Result image to write into - -> BoolReader s () -decodeImage img compCount outImage = do - zigZagArray <- lift $ createEmptyMutableMacroBlock - dcArray <- lift (M.replicate compCount 0 :: ST s (M.STVector s DcCoefficient)) - - let huffmans = gatherHuffmanTables img - huffmanForComponent dcOrAc dest = - case [t | (h,t) <- huffmans - , huffmanTableClass h == dcOrAc - , huffmanTableDest h == dest] of - (v:_) -> v - [] -> error "No Huffman table" - - mcuBeforeRestart = case [i | JpgIntervalRestart i <- jpgFrame img] of - [] -> maxBound :: Int -- HUUUUUUGE value (enough to parse all MCU) - (x:_) -> fromIntegral x - - quants = gatherQuantTables img - quantForComponent dest = - case [quantTable q | q <- quants, quantDestination q == dest] of - (v:_) -> v - [] -> error "No quant table" - - hdr = case [h | JpgScanBlob h _ <- jpgFrame img] of - [] -> error "No scan blob" - (v:_) -> v - - (_, scanInfo) = gatherScanInfo img - imgWidth = fromIntegral $ jpgWidth scanInfo - imgHeight = fromIntegral $ jpgHeight scanInfo - - blockSizeOfDim fullDim maxBlockSize = block + (if rest /= 0 then 1 else 0) - where (block, rest) = fullDim `divMod` maxBlockSize - - horizontalSamplings = [horiz | (horiz, _, _, _, _, _) <- componentsInfo] - - imgComponentCount = fromIntegral $ jpgImageComponentCount scanInfo :: Int - isImageLumanOnly = imgComponentCount == 1 - maxHorizFactor | not isImageLumanOnly && - not (allElementsEqual horizontalSamplings) = maximum horizontalSamplings - | otherwise = 1 - - verticalSamplings = [vert | (_, vert, _, _, _, _) <- componentsInfo] - maxVertFactor | not isImageLumanOnly && - not (allElementsEqual verticalSamplings) = maximum verticalSamplings - | otherwise = 1 - - horizontalBlockCount = - blockSizeOfDim imgWidth $ fromIntegral (maxHorizFactor * 8) - - verticalBlockCount = - blockSizeOfDim imgHeight $ fromIntegral (maxVertFactor * 8) - - fetchTablesForComponent component = (horizCount, vertCount, dcTree, acTree, qTable, unpacker) - where idx = componentIdentifier component - descr = case [c | c <- scans hdr, componentSelector c == idx] of - (v:_) -> v - [] -> error "No scan" - dcTree = huffmanForComponent DcComponent $ dcEntropyCodingTable descr - acTree = huffmanForComponent AcComponent $ acEntropyCodingTable descr - qTable = quantForComponent $ if idx == 1 then 0 else 1 - horizCount = if not isImageLumanOnly - then fromIntegral $ horizontalSamplingFactor component - else 1 - vertCount = if not isImageLumanOnly - then fromIntegral $ verticalSamplingFactor component - else 1 - - xScalingFactor = maxHorizFactor - horizCount + 1 - yScalingFactor = maxVertFactor - vertCount + 1 - - unpacker = unpackerDecision xScalingFactor yScalingFactor - - unpackerDecision 1 1 | isImageLumanOnly = unpack444Y - | otherwise = unpack444Ycbcr . fromIntegral $ idx - 1 - unpackerDecision 2 1 = unpack422Ycbcr . fromIntegral $ idx - 1 - unpackerDecision _ _ = unpackMacroBlock compCount (fromIntegral $ idx - 1) xScalingFactor yScalingFactor - - componentsInfo = map fetchTablesForComponent $ jpgComponents scanInfo - - let blockIndices = [(x,y) | y <- [0 .. verticalBlockCount - 1] - , x <- [0 .. horizontalBlockCount - 1] ] - blockBeforeRestart = mcuBeforeRestart - - folder f = foldM_ f blockBeforeRestart blockIndices - - folder (\resetCounter (x,y) -> do - when (resetCounter == 0) - (do forM_ [0.. compCount - 1] $ - \c -> lift $ (dcArray `M.unsafeWrite` c) 0 - byteAlignJpg - _restartCode <- decodeRestartInterval - -- if 0xD0 <= restartCode && restartCode <= 0xD7 - return ()) - - let comp _ [] = return () - comp compIdx ((horizCount, vertCount, dcTree, acTree, qTable, unpack):comp_rest) = liner 0 - where liner yd | yd >= vertCount = comp (compIdx + 1) comp_rest - liner yd = columner 0 - where verticalLimited = yd == horizCount - 1 || y == verticalBlockCount - 1 - columner xd | xd >= horizCount = liner (yd + 1) - columner xd | (xd == horizCount - 1 && x == horizontalBlockCount - 1) || verticalLimited = do - dc <- lift $ dcArray `M.unsafeRead` compIdx - (dcCoeff, block) <- - decompressMacroBlock dcTree acTree qTable zigZagArray $ fromIntegral dc - lift $ unpackMacroBlock imgComponentCount compIdx (maxHorizFactor - horizCount + 1) - (maxVertFactor - vertCount + 1) - (x * horizCount + xd) (y * vertCount + yd) outImage block - lift $ (dcArray `M.unsafeWrite` compIdx) dcCoeff - columner $ xd + 1 - - columner xd = do - dc <- lift $ dcArray `M.unsafeRead` compIdx - (dcCoeff, block) <- - decompressMacroBlock dcTree acTree qTable zigZagArray $ fromIntegral dc - _ <- lift $ unpack (x * horizCount + xd) (y * vertCount + yd) outImage block - lift $ (dcArray `M.unsafeWrite` compIdx) dcCoeff - columner $ xd + 1 - comp 0 componentsInfo - - if resetCounter /= 0 then return $ resetCounter - 1 - -- we use blockBeforeRestart - 1 to count - -- the current MCU - else return $ blockBeforeRestart - 1) - -allElementsEqual :: (Eq a) => [a] -> Bool -allElementsEqual [] = True -allElementsEqual (x:xs) = all (== x) xs - -gatherImageKind :: [JpgFrame] -> Maybe JpgImageKind -gatherImageKind lst = case [k | JpgScans k _ <- lst, isDctSpecifier k] of - [JpgBaselineDCTHuffman] -> Just BaseLineDCT - [JpgProgressiveDCTHuffman] -> Just ProgressiveDCT - _ -> Nothing - where isDctSpecifier JpgProgressiveDCTHuffman = True - isDctSpecifier JpgBaselineDCTHuffman = True - isDctSpecifier _ = False - - --- | Try to decompress a jpeg file and decompress. The colorspace is still --- YCbCr if you want to perform computation on the luma part. You can --- convert it to RGB using 'convertImage' from the 'ColorSpaceConvertible' --- typeclass. --- --- This function can output the following pixel types : --- --- * PixelY8 --- --- * PixelYCbCr8 --- -decodeJpeg :: B.ByteString -> Either String DynamicImage -decodeJpeg file = case runGetStrict get file of - Left err -> Left err - Right img -> case (compCount, imgKind) of - (_, Nothing) -> Left "Unknown Jpg kind" - (_, Just ProgressiveDCT) -> Left "Unsupported Progressive JPEG image" - (1, _) -> Right . ImageY8 $ Image imgWidth imgHeight pixelData - (3, _) -> Right . ImageYCbCr8 $ Image imgWidth imgHeight pixelData - _ -> Left "Wrong component count" - - where (imgData:_) = [d | JpgScanBlob _kind d <- jpgFrame img] - (_, scanInfo) = gatherScanInfo img - compCount = length $ jpgComponents scanInfo - - imgKind = gatherImageKind $ jpgFrame img - imgWidth = fromIntegral $ jpgWidth scanInfo - imgHeight = fromIntegral $ jpgHeight scanInfo - - imageSize = imgWidth * imgHeight * compCount - - pixelData = runST $ VS.unsafeFreeze =<< S.evalStateT (do - resultImage <- lift $ M.new imageSize - let wrapped = MutableImage imgWidth imgHeight resultImage - setDecodedStringJpg . B.concat $ L.toChunks imgData - decodeImage img compCount wrapped - return resultImage) (BoolState (-1) 0 B.empty) - -extractBlock :: Image PixelYCbCr8 -- ^ Source image - -> MutableMacroBlock s Int16 -- ^ Mutable block where to put extracted block - -> Int -- ^ Plane - -> Int -- ^ X sampling factor - -> Int -- ^ Y sampling factor - -> Int -- ^ Sample per pixel - -> Int -- ^ Block x - -> Int -- ^ Block y - -> ST s (MutableMacroBlock s Int16) -extractBlock (Image { imageWidth = w, imageHeight = h, imageData = src }) - block 1 1 sampCount plane bx by | (bx * 8) + 7 < w && (by * 8) + 7 < h = do - let baseReadIdx = (by * 8 * w) + bx * 8 - sequence_ [(block `M.unsafeWrite` (y * 8 + x)) val - | y <- [0 .. 7] - , let blockReadIdx = baseReadIdx + y * w - , x <- [0 .. 7] - , let val = fromIntegral $ src `VS.unsafeIndex` ((blockReadIdx + x) * sampCount + plane) - ] - return block -extractBlock (Image { imageWidth = w, imageHeight = h, imageData = src }) - block sampWidth sampHeight sampCount plane bx by = do - let accessPixel x y | x < w && y < h = let idx = (y * w + x) * sampCount + plane in src `VS.unsafeIndex` idx - | x >= w = accessPixel (w - 1) y - | otherwise = accessPixel x (h - 1) - - pixelPerCoeff = fromIntegral $ sampWidth * sampHeight - - blockVal x y = sum [fromIntegral $ accessPixel (xBase + dx) (yBase + dy) - | dy <- [0 .. sampHeight - 1] - , dx <- [0 .. sampWidth - 1] ] `div` pixelPerCoeff - where xBase = blockXBegin + x * sampWidth - yBase = blockYBegin + y * sampHeight - - blockXBegin = bx * 8 * sampWidth - blockYBegin = by * 8 * sampHeight - - sequence_ [(block `M.unsafeWrite` (y * 8 + x)) $ blockVal x y | y <- [0 .. 7], x <- [0 .. 7] ] - return block - -serializeMacroBlock :: BoolWriteStateRef s - -> HuffmanWriterCode -> HuffmanWriterCode - -> MutableMacroBlock s Int32 - -> ST s () -serializeMacroBlock !st !dcCode !acCode !blk = - (blk `M.unsafeRead` 0) >>= (fromIntegral >>> encodeDc) >> writeAcs (0, 1) >> return () - where writeAcs acc@(_, 63) = - (blk `M.unsafeRead` 63) >>= (fromIntegral >>> encodeAcCoefs acc) >> return () - writeAcs acc@(_, i ) = - (blk `M.unsafeRead` i) >>= (fromIntegral >>> encodeAcCoefs acc) >>= writeAcs - - encodeDc n = writeBits' st (fromIntegral code) (fromIntegral bitCount) - >> when (ssss /= 0) (encodeInt st ssss n) - where ssss = powerOf $ fromIntegral n - (bitCount, code) = dcCode `V.unsafeIndex` fromIntegral ssss - - encodeAc 0 0 = writeBits' st (fromIntegral code) $ fromIntegral bitCount - where (bitCount, code) = acCode `V.unsafeIndex` 0 - - encodeAc zeroCount n | zeroCount >= 16 = - writeBits' st (fromIntegral code) (fromIntegral bitCount) >> encodeAc (zeroCount - 16) n - where (bitCount, code) = acCode `V.unsafeIndex` 0xF0 - encodeAc zeroCount n = - writeBits' st (fromIntegral code) (fromIntegral bitCount) >> encodeInt st ssss n - where rrrr = zeroCount `unsafeShiftL` 4 - ssss = powerOf $ fromIntegral n - rrrrssss = rrrr .|. ssss - (bitCount, code) = acCode `V.unsafeIndex` fromIntegral rrrrssss - - encodeAcCoefs ( _, 63) 0 = encodeAc 0 0 >> return (0, 64) - encodeAcCoefs (zeroRunLength, i) 0 = return (zeroRunLength + 1, i + 1) - encodeAcCoefs (zeroRunLength, i) n = - encodeAc zeroRunLength n >> return (0, i + 1) - -encodeMacroBlock :: QuantificationTable - -> MutableMacroBlock s Int32 - -> MutableMacroBlock s Int32 - -> Int16 - -> MutableMacroBlock s Int16 - -> ST s (Int32, MutableMacroBlock s Int32) -encodeMacroBlock quantTableOfComponent workData finalData prev_dc block = do - -- the inverse level shift is performed internally by the fastDCT routine - blk <- fastDctLibJpeg workData block - >>= zigZagReorderForward finalData - >>= quantize quantTableOfComponent - dc <- blk `M.unsafeRead` 0 - (blk `M.unsafeWrite` 0) $ dc - fromIntegral prev_dc - return (dc, blk) - -divUpward :: (Integral a) => a -> a -> a -divUpward n dividor = val + (if rest /= 0 then 1 else 0) - where (val, rest) = n `divMod` dividor - -prepareHuffmanTable :: DctComponent -> Word8 -> HuffmanTable - -> (JpgHuffmanTableSpec, HuffmanTreeInfo) -prepareHuffmanTable classVal dest tableDef = - (JpgHuffmanTableSpec { huffmanTableClass = classVal - , huffmanTableDest = dest - , huffSizes = sizes - , huffCodes = V.fromListN 16 - [VU.fromListN (fromIntegral $ sizes ! i) lst - | (i, lst) <- zip [0..] tableDef ] - }, VS.singleton 0) - where sizes = VU.fromListN 16 $ map (fromIntegral . length) tableDef - --- | Encode an image in jpeg at a reasonnable quality level. --- If you want better quality or reduced file size, you should --- use `encodeJpegAtQuality` -encodeJpeg :: Image PixelYCbCr8 -> L.ByteString -encodeJpeg = encodeJpegAtQuality 50 - -defaultHuffmanTables :: [(JpgHuffmanTableSpec, HuffmanTreeInfo)] -defaultHuffmanTables = - [ prepareHuffmanTable DcComponent 0 defaultDcLumaHuffmanTable - , prepareHuffmanTable AcComponent 0 defaultAcLumaHuffmanTable - , prepareHuffmanTable DcComponent 1 defaultDcChromaHuffmanTable - , prepareHuffmanTable AcComponent 1 defaultAcChromaHuffmanTable - ] - --- | Function to call to encode an image to jpeg. --- The quality factor should be between 0 and 100 (100 being --- the best quality). -encodeJpegAtQuality :: Word8 -- ^ Quality factor - -> Image PixelYCbCr8 -- ^ Image to encode - -> L.ByteString -- ^ Encoded JPEG -encodeJpegAtQuality quality img@(Image { imageWidth = w, imageHeight = h }) = encode finalImage - where finalImage = JpgImage [ JpgQuantTable quantTables - , JpgScans JpgBaselineDCTHuffman hdr - , JpgHuffmanTable defaultHuffmanTables - , JpgScanBlob scanHeader encodedImage - ] - - outputComponentCount = 3 - - scanHeader = scanHeader'{ scanLength = fromIntegral $ calculateSize scanHeader' } - scanHeader' = JpgScanHeader - { scanLength = 0 - , scanComponentCount = outputComponentCount - , scans = [ JpgScanSpecification { componentSelector = 1 - , dcEntropyCodingTable = 0 - , acEntropyCodingTable = 0 - } - , JpgScanSpecification { componentSelector = 2 - , dcEntropyCodingTable = 1 - , acEntropyCodingTable = 1 - } - , JpgScanSpecification { componentSelector = 3 - , dcEntropyCodingTable = 1 - , acEntropyCodingTable = 1 - } - ] - - , spectralSelection = (0, 63) - , successiveApproxHigh = 0 - , successiveApproxLow = 0 - } - - hdr = hdr' { jpgFrameHeaderLength = fromIntegral $ calculateSize hdr' } - hdr' = JpgFrameHeader { jpgFrameHeaderLength = 0 - , jpgSamplePrecision = 8 - , jpgHeight = fromIntegral h - , jpgWidth = fromIntegral w - , jpgImageComponentCount = outputComponentCount - , jpgComponents = [ - JpgComponent { componentIdentifier = 1 - , horizontalSamplingFactor = 2 - , verticalSamplingFactor = 2 - , quantizationTableDest = 0 - } - , JpgComponent { componentIdentifier = 2 - , horizontalSamplingFactor = 1 - , verticalSamplingFactor = 1 - , quantizationTableDest = 1 - } - , JpgComponent { componentIdentifier = 3 - , horizontalSamplingFactor = 1 - , verticalSamplingFactor = 1 - , quantizationTableDest = 1 - } - ] - } - - lumaQuant = scaleQuantisationMatrix (fromIntegral quality) - defaultLumaQuantizationTable - chromaQuant = scaleQuantisationMatrix (fromIntegral quality) - defaultChromaQuantizationTable - - zigzagedLumaQuant = zigZagReorderForwardv lumaQuant - zigzagedChromaQuant = zigZagReorderForwardv chromaQuant - quantTables = [ JpgQuantTableSpec { quantPrecision = 0, quantDestination = 0 - , quantTable = zigzagedLumaQuant } - , JpgQuantTableSpec { quantPrecision = 0, quantDestination = 1 - , quantTable = zigzagedChromaQuant } - ] - - encodedImage = runST $ do - let horizontalMetaBlockCount = w `divUpward` (8 * maxSampling) - verticalMetaBlockCount = h `divUpward` (8 * maxSampling) - maxSampling = 2 - lumaSamplingSize = ( maxSampling, maxSampling, zigzagedLumaQuant - , makeInverseTable defaultDcLumaHuffmanTree - , makeInverseTable defaultAcLumaHuffmanTree) - chromaSamplingSize = ( maxSampling - 1, maxSampling - 1, zigzagedChromaQuant - , makeInverseTable defaultDcChromaHuffmanTree - , makeInverseTable defaultAcChromaHuffmanTree) - componentDef = [lumaSamplingSize, chromaSamplingSize, chromaSamplingSize] - - imageComponentCount = length componentDef - - dc_table <- M.replicate 3 0 - block <- createEmptyMutableMacroBlock - workData <- createEmptyMutableMacroBlock - zigzaged <- createEmptyMutableMacroBlock - writeState <- newWriteStateRef - - -- It's ugly, I know, be avoid allocation - let blockLine my | my >= verticalMetaBlockCount = return () - blockLine my = blockColumn 0 - where blockColumn mx | mx >= horizontalMetaBlockCount = blockLine (my + 1) - blockColumn mx = component $ zip [0..] componentDef - where component [] = blockColumn (mx + 1) - component ((comp, (sizeX, sizeY, table, dc, ac)) : comp_rest) = line 0 - where xSamplingFactor = maxSampling - sizeX + 1 - ySamplingFactor = maxSampling - sizeY + 1 - extractor = extractBlock img block xSamplingFactor ySamplingFactor imageComponentCount - line subY | subY >= sizeY = component comp_rest - line subY = column 0 - where blockY = my * sizeY + subY - - column subX | subX >= sizeX = line (subY + 1) - column subX = do - let blockX = mx * sizeX + subX - prev_dc <- dc_table `M.unsafeRead` comp - (dc_coeff, neo_block) <- (extractor comp blockX blockY >>= - encodeMacroBlock table workData zigzaged prev_dc) - (dc_table `M.unsafeWrite` comp) $ fromIntegral dc_coeff - serializeMacroBlock writeState dc ac neo_block - column $ subX + 1 - blockLine 0 - - finalizeBoolWriter writeState +module Codec.Picture.Jpg( decodeJpeg + , encodeJpegAtQuality + , encodeJpeg + + , jpgMachineStep + , JpgDecoderState( JpgDecoderState ) + ) where + +import Control.Arrow( (>>>) ) +import Control.Applicative( pure, (<$>) ) +import Control.Monad( when, forM_ ) +import Control.Monad.ST( ST, runST ) +import Control.Monad.Trans( lift ) +import Control.Monad.Trans.RWS.Strict( RWS, modify, tell, gets, execRWS ) + +import Data.Bits( (.|.), unsafeShiftL ) +import Data.Int( Int16, Int32 ) +import Data.Word(Word8, Word32) +import Data.Binary( Binary(..), encode ) +import Data.STRef( newSTRef, writeSTRef, readSTRef ) + +import Data.Vector( (//) ) +import Data.Vector.Unboxed( (!) ) +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Storable.Mutable as M +import qualified Data.ByteString as B +import qualified Data.ByteString.Lazy as L + +import Codec.Picture.InternalHelper +import Codec.Picture.BitWriter +import Codec.Picture.Types +import Codec.Picture.Jpg.Types +import Codec.Picture.Jpg.Common +import Codec.Picture.Jpg.Progressive +import Codec.Picture.Jpg.DefaultTable +import Codec.Picture.Jpg.FastDct + +quantize :: MacroBlock Int16 -> MutableMacroBlock s Int32 + -> ST s (MutableMacroBlock s Int32) +quantize table block = update 0 + where update 64 = return block + update idx = do + val <- block `M.unsafeRead` idx + let q = fromIntegral (table `VS.unsafeIndex` idx) + finalValue = (val + (q `div` 2)) `quot` q -- rounded integer division + (block `M.unsafeWrite` idx) finalValue + update $ idx + 1 + + +powerOf :: Int32 -> Word32 +powerOf 0 = 0 +powerOf n = limit 1 0 + where val = abs n + limit range i | val < range = i + limit range i = limit (2 * range) (i + 1) + +encodeInt :: BoolWriteStateRef s -> Word32 -> Int32 -> ST s () +{-# INLINE encodeInt #-} +encodeInt st ssss n | n > 0 = writeBits' st (fromIntegral n) (fromIntegral ssss) +encodeInt st ssss n = writeBits' st (fromIntegral $ n - 1) (fromIntegral ssss) + +-- | Assume the macro block is initialized with zeroes +acCoefficientsDecode :: HuffmanPackedTree -> MutableMacroBlock s Int16 + -> BoolReader s (MutableMacroBlock s Int16) +acCoefficientsDecode acTree mutableBlock = parseAcCoefficient 1 >> return mutableBlock + where parseAcCoefficient n | n >= 64 = return () + | otherwise = do + rrrrssss <- decodeRrrrSsss acTree + case rrrrssss of + ( 0, 0) -> return () + (0xF, 0) -> parseAcCoefficient (n + 16) + (rrrr, ssss) -> do + decoded <- fromIntegral <$> decodeInt ssss + lift $ (mutableBlock `M.unsafeWrite` (n + rrrr)) decoded + parseAcCoefficient (n + rrrr + 1) + +-- | Decompress a macroblock from a bitstream given the current configuration +-- from the frame. +decompressMacroBlock :: HuffmanPackedTree -- ^ Tree used for DC coefficient + -> HuffmanPackedTree -- ^ Tree used for Ac coefficient + -> MacroBlock Int16 -- ^ Current quantization table + -> MutableMacroBlock s Int16 -- ^ A zigzag table, to avoid allocation + -> DcCoefficient -- ^ Previous dc value + -> BoolReader s (DcCoefficient, MutableMacroBlock s Int16) +decompressMacroBlock dcTree acTree quantizationTable zigzagBlock previousDc = do + dcDeltaCoefficient <- dcCoefficientDecode dcTree + block <- lift createEmptyMutableMacroBlock + let neoDcCoefficient = previousDc + dcDeltaCoefficient + lift $ (block `M.unsafeWrite` 0) neoDcCoefficient + fullBlock <- acCoefficientsDecode acTree block + decodedBlock <- lift $ decodeMacroBlock quantizationTable zigzagBlock fullBlock + return (neoDcCoefficient, decodedBlock) + +pixelClamp :: Int16 -> Word8 +pixelClamp n = fromIntegral . min 255 $ max 0 n + +unpack444Y :: Int -- ^ component index + -> Int -- ^ x + -> Int -- ^ y + -> MutableImage s PixelYCbCr8 + -> MutableMacroBlock s Int16 + -> ST s () +unpack444Y _ x y (MutableImage { mutableImageWidth = imgWidth, mutableImageData = img }) + block = blockVert baseIdx 0 zero + where zero = 0 :: Int + baseIdx = x * dctBlockSize + y * dctBlockSize * imgWidth + + blockVert _ _ j | j >= dctBlockSize = return () + blockVert writeIdx readingIdx j = blockHoriz writeIdx readingIdx zero + where blockHoriz _ readIdx i | i >= dctBlockSize = blockVert (writeIdx + imgWidth) readIdx $ j + 1 + blockHoriz idx readIdx i = do + val <- pixelClamp <$> (block `M.unsafeRead` readIdx) + (img `M.unsafeWrite` idx) val + blockHoriz (idx + 1) (readIdx + 1) $ i + 1 + +unpack444Ycbcr :: Int -- ^ Component index + -> Int -- ^ x + -> Int -- ^ y + -> MutableImage s PixelYCbCr8 + -> MutableMacroBlock s Int16 + -> ST s () +unpack444Ycbcr compIdx x y + (MutableImage { mutableImageWidth = imgWidth, mutableImageData = img }) + block = blockVert baseIdx 0 zero + where zero = 0 :: Int + baseIdx = (x * dctBlockSize + y * dctBlockSize * imgWidth) * 3 + compIdx + + blockVert _ _ j | j >= dctBlockSize = return () + blockVert idx readIdx j = do + val0 <- pixelClamp <$> (block `M.unsafeRead` readIdx) + val1 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 1)) + val2 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 2)) + val3 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 3)) + val4 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 4)) + val5 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 5)) + val6 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 6)) + val7 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 7)) + + (img `M.unsafeWrite` idx) val0 + (img `M.unsafeWrite` (idx + (3 * 1))) val1 + (img `M.unsafeWrite` (idx + (3 * 2))) val2 + (img `M.unsafeWrite` (idx + (3 * 3))) val3 + (img `M.unsafeWrite` (idx + (3 * 4))) val4 + (img `M.unsafeWrite` (idx + (3 * 5))) val5 + (img `M.unsafeWrite` (idx + (3 * 6))) val6 + (img `M.unsafeWrite` (idx + (3 * 7))) val7 + + blockVert (idx + 3 * imgWidth) (readIdx + dctBlockSize) $ j + 1 + + + {-where blockHoriz _ readIdx i | i >= 8 = blockVert (writeIdx + imgWidth * 3) readIdx $ j + 1-} + {-blockHoriz idx readIdx i = do-} + {-val <- pixelClamp <$> (block `M.unsafeRead` readIdx) -} + {-(img `M.unsafeWrite` idx) val-} + {-blockHoriz (idx + 3) (readIdx + 1) $ i + 1-} + +unpack421Ycbcr :: Int -- ^ Component index + -> Int -- ^ x + -> Int -- ^ y + -> MutableImage s PixelYCbCr8 + -> MutableMacroBlock s Int16 + -> ST s () +unpack421Ycbcr compIdx x y + (MutableImage { mutableImageWidth = imgWidth, + mutableImageHeight = _, mutableImageData = img }) + block = blockVert baseIdx 0 zero + where zero = 0 :: Int + baseIdx = (x * dctBlockSize + y * dctBlockSize * imgWidth) * 3 + compIdx + lineOffset = imgWidth * 3 + + blockVert _ _ j | j >= dctBlockSize = return () + blockVert idx readIdx j = do + v0 <- pixelClamp <$> (block `M.unsafeRead` readIdx) + v1 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 1)) + v2 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 2)) + v3 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 3)) + v4 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 4)) + v5 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 5)) + v6 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 6)) + v7 <- pixelClamp <$> (block `M.unsafeRead` (readIdx + 7)) + + (img `M.unsafeWrite` idx) v0 + (img `M.unsafeWrite` (idx + 3)) v0 + + (img `M.unsafeWrite` (idx + 6 * 1)) v1 + (img `M.unsafeWrite` (idx + 6 * 1 + 3)) v1 + + (img `M.unsafeWrite` (idx + 6 * 2)) v2 + (img `M.unsafeWrite` (idx + 6 * 2 + 3)) v2 + + (img `M.unsafeWrite` (idx + 6 * 3)) v3 + (img `M.unsafeWrite` (idx + 6 * 3 + 3)) v3 + + (img `M.unsafeWrite` (idx + 6 * 4)) v4 + (img `M.unsafeWrite` (idx + 6 * 4 + 3)) v4 + + (img `M.unsafeWrite` (idx + 6 * 5)) v5 + (img `M.unsafeWrite` (idx + 6 * 5 + 3)) v5 + + (img `M.unsafeWrite` (idx + 6 * 6)) v6 + (img `M.unsafeWrite` (idx + 6 * 6 + 3)) v6 + + (img `M.unsafeWrite` (idx + 6 * 7)) v7 + (img `M.unsafeWrite` (idx + 6 * 7 + 3)) v7 + + blockVert (idx + lineOffset) (readIdx + dctBlockSize) $ j + 1 + +type Unpacker s = Int -- ^ component index + -> Int -- ^ x + -> Int -- ^ y + -> MutableImage s PixelYCbCr8 + -> MutableMacroBlock s Int16 + -> ST s () + +type JpgScripter s a = + RWS () [([(JpgUnpackerParameter, Unpacker s)], L.ByteString)] JpgDecoderState a + +data JpgDecoderState = JpgDecoderState + { dcDecoderTables :: !(V.Vector HuffmanPackedTree) + , acDecoderTables :: !(V.Vector HuffmanPackedTree) + , quantizationMatrices :: !(V.Vector (MacroBlock Int16)) + , currentRestartInterv :: !Int + , currentFrame :: Maybe JpgFrameHeader + , isProgressive :: !Bool + , maximumHorizontalResolution :: !Int + , maximumVerticalResolution :: !Int + , seenBlobs :: !Int + } + +emptyDecoderState :: JpgDecoderState +emptyDecoderState = JpgDecoderState + { dcDecoderTables = V.fromList $ map snd + [ prepareHuffmanTable DcComponent 0 defaultDcLumaHuffmanTable + , prepareHuffmanTable DcComponent 1 defaultDcChromaHuffmanTable + ] + + , acDecoderTables = V.fromList $ map snd + [ prepareHuffmanTable AcComponent 0 defaultAcLumaHuffmanTable + , prepareHuffmanTable AcComponent 1 defaultAcChromaHuffmanTable + ] + + , quantizationMatrices = V.replicate 4 (VS.replicate (8 * 8) 1) + , currentRestartInterv = -1 + , currentFrame = Nothing + , isProgressive = False + , maximumHorizontalResolution = 0 + , maximumVerticalResolution = 0 + , seenBlobs = 0 + } + +-- | This pseudo interpreter interpret the Jpg frame for the huffman, +-- quant table and restart interval parameters. +jpgMachineStep :: JpgFrame -> JpgScripter s () +jpgMachineStep (JpgAppFrame _ _) = pure () +jpgMachineStep (JpgExtension _ _) = pure () +jpgMachineStep (JpgScanBlob hdr raw_data) = do + let scanCount = length $ scans hdr + params <- concat <$> mapM (scanSpecifier scanCount) (scans hdr) + modify $ \st -> st { seenBlobs = seenBlobs st + 1 } + tell [(params, raw_data) ] + where (selectionLow, selectionHigh) = spectralSelection hdr + approxHigh = fromIntegral $ successiveApproxHigh hdr + approxLow = fromIntegral $ successiveApproxLow hdr + + scanSpecifier scanCount scanSpec = do + let dcIndex = fromIntegral $ dcEntropyCodingTable scanSpec + acIndex = fromIntegral $ acEntropyCodingTable scanSpec + comp = fromIntegral (componentSelector scanSpec) - 1 + dcTree <- gets $ (V.! dcIndex) . dcDecoderTables + acTree <- gets $ (V.! acIndex) . acDecoderTables + isProgressiveImage <- gets isProgressive + maxiW <- gets maximumHorizontalResolution + maxiH <- gets maximumVerticalResolution + restart <- gets currentRestartInterv + frameInfo <- gets currentFrame + blobId <- gets seenBlobs + case frameInfo of + Nothing -> fail "Jpg decoding error - no previous frame" + Just v -> do + let compDesc = jpgComponents v !! comp + compCount = length $ jpgComponents v + xSampling = fromIntegral $ horizontalSamplingFactor compDesc + ySampling = fromIntegral $ verticalSamplingFactor compDesc + componentSubSampling = + (maxiW - xSampling + 1, maxiH - ySampling + 1) + (xCount, yCount) + | scanCount > 1 || isProgressiveImage = (xSampling, ySampling) + | otherwise = (1, 1) + + pure [ (JpgUnpackerParameter + { dcHuffmanTree = dcTree + , acHuffmanTree = acTree + , componentIndex = comp + , restartInterval = fromIntegral restart + , componentWidth = xSampling + , componentHeight = ySampling + , subSampling = componentSubSampling + , successiveApprox = (approxLow, approxHigh) + , readerIndex = blobId + , indiceVector = + if selectionLow == 0 then 0 else 1 + , coefficientRange = + ( fromIntegral selectionLow + , fromIntegral selectionHigh ) + , blockIndex = y * ySampling + x + , blockMcuX = x + , blockMcuY = y + }, unpackerDecision compCount componentSubSampling) + | y <- [0 .. yCount - 1] + , x <- [0 .. xCount - 1] ] + +jpgMachineStep (JpgScans kind hdr) = modify $ \s -> + s { currentFrame = Just hdr + , isProgressive = case kind of + JpgProgressiveDCTHuffman -> True + _ -> False + , maximumHorizontalResolution = + fromIntegral $ maximum horizontalResolutions + , maximumVerticalResolution = + fromIntegral $ maximum verticalResolutions + } + where components = jpgComponents hdr + horizontalResolutions = map horizontalSamplingFactor components + verticalResolutions = map verticalSamplingFactor components +jpgMachineStep (JpgIntervalRestart restart) = + modify $ \s -> s { currentRestartInterv = fromIntegral restart } +jpgMachineStep (JpgHuffmanTable tables) = mapM_ placeHuffmanTrees tables + where placeHuffmanTrees (spec, tree) = case huffmanTableClass spec of + DcComponent -> modify $ \s -> + let neu = dcDecoderTables s // [(idx, tree)] in + s { dcDecoderTables = neu `seq` neu } + where idx = fromIntegral $ huffmanTableDest spec + + AcComponent -> modify $ \s -> + s { acDecoderTables = acDecoderTables s // [(idx, tree)] } + where idx = fromIntegral $ huffmanTableDest spec + +jpgMachineStep (JpgQuantTable tables) = mapM_ placeQuantizationTables tables + where placeQuantizationTables table = do + let idx = fromIntegral $ quantDestination table + tableData = quantTable table + modify $ \s -> + s { quantizationMatrices = quantizationMatrices s // [(idx, tableData)] } + +unpackerDecision :: Int -> (Int, Int) -> Unpacker s +unpackerDecision 1 (1, 1) = unpack444Y +unpackerDecision _ (1, 1) = unpack444Ycbcr +unpackerDecision _ (2, 1) = unpack421Ycbcr +unpackerDecision compCount (xScalingFactor, yScalingFactor) = + unpackMacroBlock compCount xScalingFactor yScalingFactor + +decodeImage :: JpgFrameHeader + -> V.Vector (MacroBlock Int16) + -> [([(JpgUnpackerParameter, Unpacker s)], L.ByteString)] + -> MutableImage s PixelYCbCr8 -- ^ Result image to write into + -> ST s () +decodeImage frame quants lst outImage = do + let compCount = length $ jpgComponents frame + zigZagArray <- createEmptyMutableMacroBlock + dcArray <- M.replicate compCount 0 :: ST s (M.STVector s DcCoefficient) + resetCounter <- newSTRef restartIntervalValue + + forM_ lst $ \(params, str) -> do + let componentsInfo = V.fromList params + compReader = initBoolStateJpg . B.concat $ L.toChunks str + maxiW = maximum [fst $ subSampling c | (c,_) <- params] + maxiH = maximum [snd $ subSampling c | (c,_) <- params] + mcuBlockWidth = maxiW * dctBlockSize + mcuBlockHeight = maxiH * dctBlockSize + imageMcuWidth = (imgWidth + mcuBlockWidth - 1) `div` mcuBlockWidth + imageMcuHeight = (imgHeight + mcuBlockHeight - 1) `div` mcuBlockHeight + + execBoolReader compReader $ rasterMap imageMcuWidth imageMcuHeight $ \x y -> do + resetLeft <- lift $ readSTRef resetCounter + if resetLeft == 0 then do + lift $ M.set dcArray 0 + byteAlignJpg + _restartCode <- decodeRestartInterval + lift $ resetCounter `writeSTRef` (restartIntervalValue - 1) + else + lift $ resetCounter `writeSTRef` (resetLeft - 1) + + V.forM_ componentsInfo $ \(comp, unpack) -> do + let compIdx = componentIndex comp + dcTree = dcHuffmanTree comp + acTree = acHuffmanTree comp + qTable = quants V.! (min 1 compIdx) + xd = blockMcuX comp + yd = blockMcuY comp + (subX, subY) = subSampling comp + dc <- lift $ dcArray `M.unsafeRead` compIdx + (dcCoeff, block) <- + decompressMacroBlock dcTree acTree qTable zigZagArray $ fromIntegral dc + lift $ (dcArray `M.unsafeWrite` compIdx) dcCoeff + let verticalLimited = y == imageMcuHeight - 1 + if (x == imageMcuWidth - 1) || verticalLimited then + lift $ unpackMacroBlock imgComponentCount + subX subY compIdx + (x * maxiW + xd) (y * maxiH + yd) outImage block + else + lift $ unpack compIdx (x * maxiW + xd) (y * maxiH + yd) outImage block + + where imgComponentCount = length $ jpgComponents frame + + imgWidth = fromIntegral $ jpgWidth frame + imgHeight = fromIntegral $ jpgHeight frame + restartIntervalValue = case lst of + ((p,_):_,_): _ -> restartInterval p + _ -> -1 + +gatherImageKind :: [JpgFrame] -> Maybe JpgImageKind +gatherImageKind lst = case [k | JpgScans k _ <- lst, isDctSpecifier k] of + [JpgBaselineDCTHuffman] -> Just BaseLineDCT + [JpgProgressiveDCTHuffman] -> Just ProgressiveDCT + _ -> Nothing + where isDctSpecifier JpgProgressiveDCTHuffman = True + isDctSpecifier JpgBaselineDCTHuffman = True + isDctSpecifier _ = False + +gatherScanInfo :: JpgImage -> (JpgFrameKind, JpgFrameHeader) +gatherScanInfo img = head [(a, b) | JpgScans a b <- jpgFrame img] + +-- | Try to decompress a jpeg file and decompress. The colorspace is still +-- YCbCr if you want to perform computation on the luma part. You can +-- convert it to RGB using 'convertImage' from the 'ColorSpaceConvertible' +-- typeclass. +-- +-- This function can output the following pixel types : +-- +-- * PixelY8 +-- +-- * PixelYCbCr8 +-- +decodeJpeg :: B.ByteString -> Either String DynamicImage +decodeJpeg file = case runGetStrict get file of + Left err -> Left err + Right img -> case (compCount, imgKind) of + (_, Nothing) -> Left "Unknown Jpg kind" + (3, Just ProgressiveDCT) -> Right . ImageYCbCr8 $ decodeProgressive + (1, Just BaseLineDCT) -> Right . ImageY8 $ Image imgWidth imgHeight pixelData + (3, Just BaseLineDCT) -> Right . ImageYCbCr8 $ Image imgWidth imgHeight pixelData + _ -> Left "Wrong component count" + + where compCount = length $ jpgComponents scanInfo + (_,scanInfo) = gatherScanInfo img + + imgKind = gatherImageKind $ jpgFrame img + imgWidth = fromIntegral $ jpgWidth scanInfo + imgHeight = fromIntegral $ jpgHeight scanInfo + + imageSize = imgWidth * imgHeight * compCount + (st, wrotten) = execRWS (mapM_ jpgMachineStep (jpgFrame img)) () emptyDecoderState + Just fHdr = currentFrame st + + decodeProgressive = runST $ + progressiveUnpack + (maximumHorizontalResolution st, maximumVerticalResolution st) + fHdr + (quantizationMatrices st) + wrotten >>= unsafeFreezeImage + + pixelData = runST $ do + resultImage <- M.new imageSize + let wrapped = MutableImage imgWidth imgHeight resultImage + decodeImage + fHdr + (quantizationMatrices st) + wrotten + wrapped + VS.unsafeFreeze resultImage + +extractBlock :: Image PixelYCbCr8 -- ^ Source image + -> MutableMacroBlock s Int16 -- ^ Mutable block where to put extracted block + -> Int -- ^ Plane + -> Int -- ^ X sampling factor + -> Int -- ^ Y sampling factor + -> Int -- ^ Sample per pixel + -> Int -- ^ Block x + -> Int -- ^ Block y + -> ST s (MutableMacroBlock s Int16) +extractBlock (Image { imageWidth = w, imageHeight = h, imageData = src }) + block 1 1 sampCount plane bx by | (bx * dctBlockSize) + 7 < w && (by * 8) + 7 < h = do + let baseReadIdx = (by * dctBlockSize * w) + bx * dctBlockSize + sequence_ [(block `M.unsafeWrite` (y * dctBlockSize + x)) val + | y <- [0 .. dctBlockSize - 1] + , let blockReadIdx = baseReadIdx + y * w + , x <- [0 .. dctBlockSize - 1] + , let val = fromIntegral $ src `VS.unsafeIndex` ((blockReadIdx + x) * sampCount + plane) + ] + return block +extractBlock (Image { imageWidth = w, imageHeight = h, imageData = src }) + block sampWidth sampHeight sampCount plane bx by = do + let accessPixel x y | x < w && y < h = let idx = (y * w + x) * sampCount + plane in src `VS.unsafeIndex` idx + | x >= w = accessPixel (w - 1) y + | otherwise = accessPixel x (h - 1) + + pixelPerCoeff = fromIntegral $ sampWidth * sampHeight + + blockVal x y = sum [fromIntegral $ accessPixel (xBase + dx) (yBase + dy) + | dy <- [0 .. sampHeight - 1] + , dx <- [0 .. sampWidth - 1] ] `div` pixelPerCoeff + where xBase = blockXBegin + x * sampWidth + yBase = blockYBegin + y * sampHeight + + blockXBegin = bx * dctBlockSize * sampWidth + blockYBegin = by * dctBlockSize * sampHeight + + sequence_ [(block `M.unsafeWrite` (y * dctBlockSize + x)) $ blockVal x y | y <- [0 .. 7], x <- [0 .. 7] ] + return block + +serializeMacroBlock :: BoolWriteStateRef s + -> HuffmanWriterCode -> HuffmanWriterCode + -> MutableMacroBlock s Int32 + -> ST s () +serializeMacroBlock !st !dcCode !acCode !blk = + (blk `M.unsafeRead` 0) >>= (fromIntegral >>> encodeDc) >> writeAcs (0, 1) >> return () + where writeAcs acc@(_, 63) = + (blk `M.unsafeRead` 63) >>= (fromIntegral >>> encodeAcCoefs acc) >> return () + writeAcs acc@(_, i ) = + (blk `M.unsafeRead` i) >>= (fromIntegral >>> encodeAcCoefs acc) >>= writeAcs + + encodeDc n = writeBits' st (fromIntegral code) (fromIntegral bitCount) + >> when (ssss /= 0) (encodeInt st ssss n) + where ssss = powerOf $ fromIntegral n + (bitCount, code) = dcCode `V.unsafeIndex` fromIntegral ssss + + encodeAc 0 0 = writeBits' st (fromIntegral code) $ fromIntegral bitCount + where (bitCount, code) = acCode `V.unsafeIndex` 0 + + encodeAc zeroCount n | zeroCount >= 16 = + writeBits' st (fromIntegral code) (fromIntegral bitCount) >> encodeAc (zeroCount - 16) n + where (bitCount, code) = acCode `V.unsafeIndex` 0xF0 + encodeAc zeroCount n = + writeBits' st (fromIntegral code) (fromIntegral bitCount) >> encodeInt st ssss n + where rrrr = zeroCount `unsafeShiftL` 4 + ssss = powerOf $ fromIntegral n + rrrrssss = rrrr .|. ssss + (bitCount, code) = acCode `V.unsafeIndex` fromIntegral rrrrssss + + encodeAcCoefs ( _, 63) 0 = encodeAc 0 0 >> return (0, 64) + encodeAcCoefs (zeroRunLength, i) 0 = return (zeroRunLength + 1, i + 1) + encodeAcCoefs (zeroRunLength, i) n = + encodeAc zeroRunLength n >> return (0, i + 1) + +encodeMacroBlock :: QuantificationTable + -> MutableMacroBlock s Int32 + -> MutableMacroBlock s Int32 + -> Int16 + -> MutableMacroBlock s Int16 + -> ST s (Int32, MutableMacroBlock s Int32) +encodeMacroBlock quantTableOfComponent workData finalData prev_dc block = do + -- the inverse level shift is performed internally by the fastDCT routine + blk <- fastDctLibJpeg workData block + >>= zigZagReorderForward finalData + >>= quantize quantTableOfComponent + dc <- blk `M.unsafeRead` 0 + (blk `M.unsafeWrite` 0) $ dc - fromIntegral prev_dc + return (dc, blk) + +divUpward :: (Integral a) => a -> a -> a +divUpward n dividor = val + (if rest /= 0 then 1 else 0) + where (val, rest) = n `divMod` dividor + +prepareHuffmanTable :: DctComponent -> Word8 -> HuffmanTable + -> (JpgHuffmanTableSpec, HuffmanPackedTree) +prepareHuffmanTable classVal dest tableDef = + (JpgHuffmanTableSpec { huffmanTableClass = classVal + , huffmanTableDest = dest + , huffSizes = sizes + , huffCodes = V.fromListN 16 + [VU.fromListN (fromIntegral $ sizes ! i) lst + | (i, lst) <- zip [0..] tableDef ] + }, VS.singleton 0) + where sizes = VU.fromListN 16 $ map (fromIntegral . length) tableDef + +-- | Encode an image in jpeg at a reasonnable quality level. +-- If you want better quality or reduced file size, you should +-- use `encodeJpegAtQuality` +encodeJpeg :: Image PixelYCbCr8 -> L.ByteString +encodeJpeg = encodeJpegAtQuality 50 + +defaultHuffmanTables :: [(JpgHuffmanTableSpec, HuffmanPackedTree)] +defaultHuffmanTables = + [ prepareHuffmanTable DcComponent 0 defaultDcLumaHuffmanTable + , prepareHuffmanTable AcComponent 0 defaultAcLumaHuffmanTable + , prepareHuffmanTable DcComponent 1 defaultDcChromaHuffmanTable + , prepareHuffmanTable AcComponent 1 defaultAcChromaHuffmanTable + ] + +-- | Function to call to encode an image to jpeg. +-- The quality factor should be between 0 and 100 (100 being +-- the best quality). +encodeJpegAtQuality :: Word8 -- ^ Quality factor + -> Image PixelYCbCr8 -- ^ Image to encode + -> L.ByteString -- ^ Encoded JPEG +encodeJpegAtQuality quality img@(Image { imageWidth = w, imageHeight = h }) = encode finalImage + where finalImage = JpgImage [ JpgQuantTable quantTables + , JpgScans JpgBaselineDCTHuffman hdr + , JpgHuffmanTable defaultHuffmanTables + , JpgScanBlob scanHeader encodedImage + ] + + outputComponentCount = 3 + + scanHeader = scanHeader'{ scanLength = fromIntegral $ calculateSize scanHeader' } + scanHeader' = JpgScanHeader + { scanLength = 0 + , scanComponentCount = outputComponentCount + , scans = [ JpgScanSpecification { componentSelector = 1 + , dcEntropyCodingTable = 0 + , acEntropyCodingTable = 0 + } + , JpgScanSpecification { componentSelector = 2 + , dcEntropyCodingTable = 1 + , acEntropyCodingTable = 1 + } + , JpgScanSpecification { componentSelector = 3 + , dcEntropyCodingTable = 1 + , acEntropyCodingTable = 1 + } + ] + + , spectralSelection = (0, 63) + , successiveApproxHigh = 0 + , successiveApproxLow = 0 + } + + hdr = hdr' { jpgFrameHeaderLength = fromIntegral $ calculateSize hdr' } + hdr' = JpgFrameHeader { jpgFrameHeaderLength = 0 + , jpgSamplePrecision = 8 + , jpgHeight = fromIntegral h + , jpgWidth = fromIntegral w + , jpgImageComponentCount = outputComponentCount + , jpgComponents = [ + JpgComponent { componentIdentifier = 1 + , horizontalSamplingFactor = 2 + , verticalSamplingFactor = 2 + , quantizationTableDest = 0 + } + , JpgComponent { componentIdentifier = 2 + , horizontalSamplingFactor = 1 + , verticalSamplingFactor = 1 + , quantizationTableDest = 1 + } + , JpgComponent { componentIdentifier = 3 + , horizontalSamplingFactor = 1 + , verticalSamplingFactor = 1 + , quantizationTableDest = 1 + } + ] + } + + lumaQuant = scaleQuantisationMatrix (fromIntegral quality) + defaultLumaQuantizationTable + chromaQuant = scaleQuantisationMatrix (fromIntegral quality) + defaultChromaQuantizationTable + + zigzagedLumaQuant = zigZagReorderForwardv lumaQuant + zigzagedChromaQuant = zigZagReorderForwardv chromaQuant + quantTables = [ JpgQuantTableSpec { quantPrecision = 0, quantDestination = 0 + , quantTable = zigzagedLumaQuant } + , JpgQuantTableSpec { quantPrecision = 0, quantDestination = 1 + , quantTable = zigzagedChromaQuant } + ] + + encodedImage = runST $ do + let horizontalMetaBlockCount = + w `divUpward` (dctBlockSize * maxSampling) + verticalMetaBlockCount = + h `divUpward` (dctBlockSize * maxSampling) + maxSampling = 2 + lumaSamplingSize = ( maxSampling, maxSampling, zigzagedLumaQuant + , makeInverseTable defaultDcLumaHuffmanTree + , makeInverseTable defaultAcLumaHuffmanTree) + chromaSamplingSize = ( maxSampling - 1, maxSampling - 1, zigzagedChromaQuant + , makeInverseTable defaultDcChromaHuffmanTree + , makeInverseTable defaultAcChromaHuffmanTree) + componentDef = [lumaSamplingSize, chromaSamplingSize, chromaSamplingSize] + + imageComponentCount = length componentDef + + dc_table <- M.replicate 3 0 + block <- createEmptyMutableMacroBlock + workData <- createEmptyMutableMacroBlock + zigzaged <- createEmptyMutableMacroBlock + writeState <- newWriteStateRef + + -- It's ugly, I know, be avoid allocation + let blockDecoder mx my = component $ zip [0..] componentDef + where component [] = return () + component ((comp, (sizeX, sizeY, table, dc, ac)) : comp_rest) = + rasterMap sizeX sizeY decoder >> component comp_rest + where xSamplingFactor = maxSampling - sizeX + 1 + ySamplingFactor = maxSampling - sizeY + 1 + extractor = extractBlock img block xSamplingFactor ySamplingFactor imageComponentCount + + decoder subX subY = do + let blockY = my * sizeY + subY + blockX = mx * sizeX + subX + prev_dc <- dc_table `M.unsafeRead` comp + (dc_coeff, neo_block) <- (extractor comp blockX blockY >>= + encodeMacroBlock table workData zigzaged prev_dc) + (dc_table `M.unsafeWrite` comp) $ fromIntegral dc_coeff + serializeMacroBlock writeState dc ac neo_block + + rasterMap + horizontalMetaBlockCount verticalMetaBlockCount + blockDecoder + + finalizeBoolWriter writeState + diff --git a/Codec/Picture/Jpg/Common.hs b/Codec/Picture/Jpg/Common.hs new file mode 100644 --- /dev/null +++ b/Codec/Picture/Jpg/Common.hs @@ -0,0 +1,232 @@ +{-# LANGUAGE BangPatterns #-} +module Codec.Picture.Jpg.Common + ( DctCoefficients + , JpgUnpackerParameter( .. ) + , decodeInt + , dcCoefficientDecode + , deQuantize + , decodeRrrrSsss + , zigZagReorderForward + , zigZagReorderForwardv + , zigZagReorder + , inverseDirectCosineTransform + , unpackInt + , unpackMacroBlock + , rasterMap + , decodeMacroBlock + , decodeRestartInterval + ) where + +import Control.Applicative( (<$>), pure ) +import Control.Monad( replicateM, when ) +import Control.Monad.ST( ST, runST ) +import Data.Bits( unsafeShiftL, unsafeShiftR, (.&.) ) +import Data.Int( Int16, Int32 ) +import Data.List( foldl' ) +import Data.Maybe( fromMaybe ) +import Data.Word( Word8 ) +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Storable.Mutable as M +import Foreign.Storable ( Storable ) + +import Codec.Picture.Types +import Codec.Picture.BitWriter +import Codec.Picture.Jpg.Types +import Codec.Picture.Jpg.FastIdct +import Codec.Picture.Jpg.DefaultTable + +-- | Same as for DcCoefficient, to provide nicer type signatures +type DctCoefficients = DcCoefficient + +data JpgUnpackerParameter = JpgUnpackerParameter + { dcHuffmanTree :: !HuffmanPackedTree + , acHuffmanTree :: !HuffmanPackedTree + , componentIndex :: {-# UNPACK #-} !Int + , restartInterval :: {-# UNPACK #-} !Int + , componentWidth :: {-# UNPACK #-} !Int + , componentHeight :: {-# UNPACK #-} !Int + , subSampling :: !(Int, Int) + , coefficientRange :: !(Int, Int) + , successiveApprox :: !(Int, Int) + , readerIndex :: {-# UNPACK #-} !Int + , indiceVector :: {-# UNPACK #-} !Int + , blockIndex :: {-# UNPACK #-} !Int + , blockMcuX :: {-# UNPACK #-} !Int + , blockMcuY :: {-# UNPACK #-} !Int + } + deriving Show + +decodeRestartInterval :: BoolReader s Int32 +decodeRestartInterval = return (-1) {- do + bits <- replicateM 8 getNextBitJpg + if bits == replicate 8 True + then do + marker <- replicateM 8 getNextBitJpg + return $ packInt marker + else return (-1) + -} + +{-# INLINE decodeInt #-} +decodeInt :: Int -> BoolReader s Int32 +decodeInt ssss = do + signBit <- getNextBitJpg + let dataRange = 1 `unsafeShiftL` fromIntegral (ssss - 1) + leftBitCount = ssss - 1 + -- First following bits store the sign of the coefficient, and counted in + -- SSSS, so the bit count for the int, is ssss - 1 + if signBit + then (\w -> dataRange + fromIntegral w) <$> unpackInt leftBitCount + else (\w -> 1 - dataRange * 2 + fromIntegral w) <$> unpackInt leftBitCount + +decodeRrrrSsss :: HuffmanPackedTree -> BoolReader s (Int, Int) +decodeRrrrSsss tree = do + rrrrssss <- huffmanPackedDecode tree + let rrrr = (rrrrssss `unsafeShiftR` 4) .&. 0xF + ssss = rrrrssss .&. 0xF + pure (fromIntegral rrrr, fromIntegral ssss) + +dcCoefficientDecode :: HuffmanPackedTree -> BoolReader s DcCoefficient +dcCoefficientDecode dcTree = do + ssss <- huffmanPackedDecode dcTree + if ssss == 0 + then return 0 + else fromIntegral <$> decodeInt (fromIntegral ssss) + +-- | Apply a quantization matrix to a macroblock +{-# INLINE deQuantize #-} +deQuantize :: MacroBlock Int16 -> MutableMacroBlock s Int16 + -> ST s (MutableMacroBlock s Int16) +deQuantize table block = update 0 + where update 64 = return block + update i = do + val <- block `M.unsafeRead` i + let finalValue = val * (table `VS.unsafeIndex` i) + (block `M.unsafeWrite` i) finalValue + update $ i + 1 + +inverseDirectCosineTransform :: MutableMacroBlock s Int16 + -> ST s (MutableMacroBlock s Int16) +inverseDirectCosineTransform mBlock = + fastIdct mBlock >>= mutableLevelShift + +zigZagOrder :: MacroBlock Int +zigZagOrder = makeMacroBlock $ concat + [[ 0, 1, 5, 6,14,15,27,28] + ,[ 2, 4, 7,13,16,26,29,42] + ,[ 3, 8,12,17,25,30,41,43] + ,[ 9,11,18,24,31,40,44,53] + ,[10,19,23,32,39,45,52,54] + ,[20,22,33,38,46,51,55,60] + ,[21,34,37,47,50,56,59,61] + ,[35,36,48,49,57,58,62,63] + ] + +zigZagReorderForwardv :: (Storable a, Num a) => VS.Vector a -> VS.Vector a +zigZagReorderForwardv vec = runST $ do + v <- M.new 64 + mv <- VS.thaw vec + zigZagReorderForward v mv >>= VS.freeze + +zigZagOrderForward :: MacroBlock Int +zigZagOrderForward = VS.generate 64 inv + where inv i = fromMaybe 0 $ VS.findIndex (i ==) zigZagOrder + +zigZagReorderForward :: (Storable a, Num a) + => MutableMacroBlock s a + -> MutableMacroBlock s a + -> ST s (MutableMacroBlock s a) +{-# SPECIALIZE INLINE zigZagReorderForward :: MutableMacroBlock s Int32 + -> MutableMacroBlock s Int32 + -> ST s (MutableMacroBlock s Int32) #-} +{-# SPECIALIZE INLINE zigZagReorderForward :: MutableMacroBlock s Int16 + -> MutableMacroBlock s Int16 + -> ST s (MutableMacroBlock s Int16) #-} +{-# SPECIALIZE INLINE zigZagReorderForward :: MutableMacroBlock s Word8 + -> MutableMacroBlock s Word8 + -> ST s (MutableMacroBlock s Word8) #-} +zigZagReorderForward zigzaged block = ordering zigZagOrderForward >> return zigzaged + where ordering !table = reorder (0 :: Int) + where reorder !i | i >= 64 = return () + reorder i = do + let idx = table `VS.unsafeIndex` i + v <- block `M.unsafeRead` idx + (zigzaged `M.unsafeWrite` i) v + reorder (i + 1) + +zigZagReorder :: MutableMacroBlock s Int16 -> MutableMacroBlock s Int16 + -> ST s (MutableMacroBlock s Int16) +zigZagReorder zigzaged block = do + let update i = do + let idx = zigZagOrder `VS.unsafeIndex` i + v <- block `M.unsafeRead` idx + (zigzaged `M.unsafeWrite` i) v + + reorder 63 = update 63 + reorder i = update i >> reorder (i + 1) + + reorder (0 :: Int) + return zigzaged + +-- | Unpack an int of the given size encoded from MSB to LSB. +unpackInt :: Int -> BoolReader s Int32 +unpackInt bitCount = packInt <$> replicateM bitCount getNextBitJpg + + +{-# INLINE rasterMap #-} +rasterMap :: (Monad m) + => Int -> Int -> (Int -> Int -> m ()) + -> m () +rasterMap width height f = liner 0 + where liner y | y >= height = return () + liner y = columner 0 + where columner x | x >= width = liner (y + 1) + columner x = f x y >> columner (x + 1) + +packInt :: [Bool] -> Int32 +packInt = foldl' bitStep 0 + where bitStep acc True = (acc `unsafeShiftL` 1) + 1 + bitStep acc False = acc `unsafeShiftL` 1 + +pixelClamp :: Int16 -> Word8 +pixelClamp n = fromIntegral . min 255 $ max 0 n + +-- | Given a size coefficient (how much a pixel span horizontally +-- and vertically), the position of the macroblock, return a list +-- of indices and value to be stored in an array (like the final +-- image) +unpackMacroBlock :: Int -- ^ Component count + -> Int -- ^ Width coefficient + -> Int -- ^ Height coefficient + -> Int -- ^ x + -> Int -- ^ y + -> Int -- ^ Component index + -> MutableImage s PixelYCbCr8 + -> MutableMacroBlock s Int16 + -> ST s () +unpackMacroBlock compCount wCoeff hCoeff compIdx x y + (MutableImage { mutableImageWidth = imgWidth, + mutableImageHeight = imgHeight, mutableImageData = img }) + block = rasterMap dctBlockSize dctBlockSize unpacker + where unpacker i j = do + let yBase = y * dctBlockSize + j * hCoeff + compVal <- pixelClamp <$> (block `M.unsafeRead` (i + j * dctBlockSize)) + rasterMap wCoeff hCoeff $ \wDup hDup -> do + let xBase = x * dctBlockSize + i * wCoeff + xPos = xBase + wDup + yPos = yBase + hDup + + when (xPos < imgWidth && yPos < imgHeight) + (do let mutableIdx = (xPos + yPos * imgWidth) * compCount + compIdx + (img `M.unsafeWrite` mutableIdx) compVal) + +-- | This is one of the most important function of the decoding, +-- it form the barebone decoding pipeline for macroblock. It's all +-- there is to know for macro block transformation +decodeMacroBlock :: MacroBlock DctCoefficients + -> MutableMacroBlock s Int16 + -> MutableMacroBlock s Int16 + -> ST s (MutableMacroBlock s Int16) +decodeMacroBlock quantizationTable zigZagBlock block = + deQuantize quantizationTable block >>= zigZagReorder zigZagBlock + >>= inverseDirectCosineTransform + diff --git a/Codec/Picture/Jpg/DefaultTable.hs b/Codec/Picture/Jpg/DefaultTable.hs --- a/Codec/Picture/Jpg/DefaultTable.hs +++ b/Codec/Picture/Jpg/DefaultTable.hs @@ -14,6 +14,7 @@ , makeInverseTable , buildHuffmanTree , packHuffmanTree + , huffmanPackedDecode , defaultChromaQuantizationTable @@ -37,11 +38,13 @@ import Control.Monad.ST( runST ) import qualified Data.Vector.Storable as SV import qualified Data.Vector as V -import Data.Bits( unsafeShiftL, (.|.) ) +import Data.Bits( unsafeShiftL, (.|.), (.&.) ) import Data.Word( Word8, Word16 ) import Data.List( foldl' ) import qualified Data.Vector.Storable.Mutable as M +import Codec.Picture.BitWriter + -- | Tree storing the code used for huffman encoding. data HuffmanTree = Branch HuffmanTree HuffmanTree -- ^ If bit is 0 take the first subtree, if 1, the right. | Leaf Word8 -- ^ We should output the value @@ -72,7 +75,7 @@ aux (Branch i1 i2@(Leaf _)) idx = do ix1 <- aux i1 (idx + 2) _ <- aux i2 (idx + 1) - (table `M.unsafeWrite` idx) (fromIntegral ix1) + (table `M.unsafeWrite` idx) . fromIntegral $ idx + 2 return ix1 aux (Branch i1 i2) idx = do @@ -137,6 +140,16 @@ scale coeff i = fromIntegral . min 255 . max 1 $ fromIntegral i * coeff `div` 100 + +huffmanPackedDecode :: HuffmanPackedTree -> BoolReader s Word8 +huffmanPackedDecode table = getNextBitJpg >>= aux 0 + where aux idx b + | (v .&. 0x8000) /= 0 = return 0 + | (v .&. 0x4000) /= 0 = return . fromIntegral $ v .&. 0xFF + | otherwise = getNextBitJpg >>= aux v + where tableIndex | b = idx + 1 + | otherwise = idx + v = table `SV.unsafeIndex` fromIntegral tableIndex defaultLumaQuantizationTable :: QuantificationTable defaultLumaQuantizationTable = makeMacroBlock diff --git a/Codec/Picture/Jpg/FastDct.hs b/Codec/Picture/Jpg/FastDct.hs --- a/Codec/Picture/Jpg/FastDct.hs +++ b/Codec/Picture/Jpg/FastDct.hs @@ -16,13 +16,13 @@ referenceDct :: MutableMacroBlock s Int32 -> MutableMacroBlock s Int16 -> ST s (MutableMacroBlock s Int32) -referenceDct workData block = forM_ [(u, v) | u <- [0 :: Int .. 7], v <- [0..7]] (\(u,v) -> do +referenceDct workData block = forM_ [(u, v) | u <- [0 :: Int .. dctBlockSize - 1], v <- [0..dctBlockSize - 1]] (\(u,v) -> do val <- at (u,v) - (workData `M.unsafeWrite` (v * 8 + u)) . truncate $ (1 / 4) * c u * c v * val) + (workData `M.unsafeWrite` (v * dctBlockSize + u)) . truncate $ (1 / 4) * c u * c v * val) >> return workData where -- at :: (Int, Int) -> ST s Float - at (u,v) = sum <$> (forM [(x,y) | x <- [0..7], y <- [0..7 :: Int]] $ \(x,y) -> do - sample <- fromIntegral <$> (block `M.unsafeRead` (y * 8 + x)) + at (u,v) = sum <$> (forM [(x,y) | x <- [0..dctBlockSize - 1], y <- [0..dctBlockSize - 1 :: Int]] $ \(x,y) -> do + sample <- fromIntegral <$> (block `M.unsafeRead` (y * dctBlockSize + x)) return $ sample * cos ((2 * fromIntegral x + 1) * fromIntegral u * (pi :: Float)/ 16) * cos ((2 * fromIntegral y + 1) * fromIntegral v * pi / 16)) c 0 = 1 / sqrt 2 @@ -65,9 +65,9 @@ where -- Pass 1: process rows. -- Note results are scaled up by sqrt(8) compared to a true DCT; -- furthermore, we scale the results by 2**PASS1_BITS. - firstPass _ 8 = return () + firstPass _ i | i == dctBlockSize = return () firstPass dataBlock i = do - let baseIdx = i * 8 + let baseIdx = i * dctBlockSize readAt idx = fromIntegral <$> sample_block `M.unsafeRead` (baseIdx + idx) mult = (*) writeAt idx n = (dataBlock `M.unsafeWrite` (baseIdx + idx)) n @@ -99,7 +99,7 @@ tmp3' = blk3 - blk4 -- Stage 4 and output - writeAt 0 $ (tmp10 + tmp11 - 8 * cENTERJSAMPLE) `unsafeShiftL` pASS1_BITS + writeAt 0 $ (tmp10 + tmp11 - dctBlockSize * cENTERJSAMPLE) `unsafeShiftL` pASS1_BITS writeAt 4 $ (tmp10 - tmp11) `unsafeShiftL` pASS1_BITS let z1 = mult (tmp12 + tmp13) fIX_0_541196100 @@ -138,10 +138,10 @@ secondPass :: M.STVector s Int32 -> Int -> ST s () secondPass _ (-1) = return () secondPass block i = do - let readAt idx = block `M.unsafeRead` ((7 - i) + idx * 8) + let readAt idx = block `M.unsafeRead` ((7 - i) + idx * dctBlockSize) mult = (*) - writeAt idx n = (block `M.unsafeWrite` (8 * idx + (7 - i))) n - writeAtPos idx n = (block `M.unsafeWrite` (8 * idx + (7 - i))) $ n `unsafeShiftR` (cONST_BITS + pASS1_BITS + 3) + writeAt idx n = (block `M.unsafeWrite` (dctBlockSize * idx + (7 - i))) n + writeAtPos idx n = (block `M.unsafeWrite` (dctBlockSize * idx + (7 - i))) $ n `unsafeShiftR` (cONST_BITS + pASS1_BITS + 3) blk0 <- readAt 0 blk1 <- readAt 1 blk2 <- readAt 2 diff --git a/Codec/Picture/Jpg/Progressive.hs b/Codec/Picture/Jpg/Progressive.hs new file mode 100644 --- /dev/null +++ b/Codec/Picture/Jpg/Progressive.hs @@ -0,0 +1,309 @@ +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeSynonymInstances #-} +{-# LANGUAGE FlexibleInstances #-} +module Codec.Picture.Jpg.Progressive + ( JpgUnpackerParameter( .. ) + , progressiveUnpack + ) where + +import Control.Applicative( pure, (<$>), (<*>) ) +import Control.Monad( when, forM_ ) +import Control.Monad.ST( ST ) +import Control.Monad.Trans( lift ) +import Data.Bits( (.&.), (.|.), unsafeShiftL ) +import Data.Int( Int16, Int32 ) +import qualified Data.ByteString as B +import qualified Data.ByteString.Lazy as L +import qualified Data.Vector as V +import qualified Data.Vector.Storable as VS +import Data.Vector( (!) ) +import qualified Data.Vector.Mutable as M +import qualified Data.Vector.Storable.Mutable as MS + +import Codec.Picture.Types +import Codec.Picture.BitWriter +import Codec.Picture.Jpg.Common +import Codec.Picture.Jpg.Types +import Codec.Picture.Jpg.DefaultTable + +mcuIndexer :: JpgComponent -> DctComponent -> Int -> VS.Vector Int +mcuIndexer param kind mcuWidth = VS.fromListN (mcuWidth * th) indexes + where compW = fromIntegral $ horizontalSamplingFactor param + compH = fromIntegral $ verticalSamplingFactor param + th = compW * compH + + index AcComponent mcu y x = mcu * th + x + y * compW + index DcComponent mcu y x = (mcu + y * mcuWidth) * compW + x + indexes = + index kind <$> [0 .. mcuWidth - 1] <*> [0 .. compH - 1] <*> [0 .. compW - 1] + +createMcuLineIndices :: JpgComponent -> Int -> V.Vector (VS.Vector Int) +createMcuLineIndices comp mcuWidth = + V.fromList [ mcuIndexer comp part mcuWidth | part <- [DcComponent, AcComponent] ] + +decodeFirstDC :: JpgUnpackerParameter + -> MS.STVector s Int16 + -> MutableMacroBlock s Int16 + -> Int32 + -> BoolReader s Int32 +decodeFirstDC params dcCoeffs block eobrun = unpack >> pure eobrun + where unpack = do + (dcDeltaCoefficient) <- dcCoefficientDecode $ dcHuffmanTree params + previousDc <- lift $ dcCoeffs `MS.unsafeRead` componentIndex params + let neoDcCoefficient = previousDc + dcDeltaCoefficient + approxLow = fst $ successiveApprox params + scaledDc = neoDcCoefficient `unsafeShiftL` approxLow + lift $ (block `MS.unsafeWrite` 0) scaledDc + lift $ (dcCoeffs `MS.unsafeWrite` componentIndex params) neoDcCoefficient + +decodeRefineDc :: JpgUnpackerParameter + -> a + -> MutableMacroBlock s Int16 + -> Int32 + -> BoolReader s Int32 +decodeRefineDc params _ block eobrun = unpack >> pure eobrun + where approxLow = fst $ successiveApprox params + plusOne = 1 `unsafeShiftL` approxLow + unpack = do + bit <- getNextBitJpg + when bit . lift $ do + v <- block `MS.unsafeRead` 0 + (block `MS.unsafeWrite` 0) $ v .|. plusOne + +decodeFirstAc :: JpgUnpackerParameter + -> a + -> MutableMacroBlock s Int16 + -> Int32 + -> BoolReader s Int32 +decodeFirstAc _params _ _block eobrun | eobrun > 0 = pure $ eobrun - 1 +decodeFirstAc params _ block _ = unpack startIndex + where (startIndex, maxIndex) = coefficientRange params + (low, _) = successiveApprox params + unpack n | n > maxIndex = pure 0 + unpack n = do + rrrrssss <- decodeRrrrSsss $ acHuffmanTree params + case rrrrssss of + (0xF, 0) -> unpack $ n + 16 + ( 0, 0) -> return 0 + ( r, 0) -> eobrun <$> unpackInt r + where eobrun lowBits = (1 `unsafeShiftL` r) - 1 + lowBits + ( r, s) -> do + let n' = n + r + val <- (`unsafeShiftL` low) <$> decodeInt s + lift . (block `MS.unsafeWrite` n') $ fromIntegral val + unpack $ n' + 1 + +decodeRefineAc :: JpgUnpackerParameter + -> a + -> MutableMacroBlock s Int16 + -> Int32 + -> BoolReader s Int32 +decodeRefineAc params _ block eobrun + | eobrun == 0 = unpack startIndex + | otherwise = performEobRun startIndex >> return (eobrun - 1) + where (startIndex, maxIndex) = coefficientRange params + (low, _) = successiveApprox params + plusOne = 1 `unsafeShiftL` low + minusOne = (-1) `unsafeShiftL` low + + getBitVal = do + v <- getNextBitJpg + pure $ if v then plusOne else minusOne + + {-performEobRun :: Int -> BoolReader s ()-} + performEobRun idx | idx > maxIndex = pure () + performEobRun idx = do + coeff <- lift $ block `MS.unsafeRead` idx + when (coeff /= 0) $ do + bit <- getNextBitJpg + case (bit, (coeff .&. plusOne) == 0) of + (False, _) -> performEobRun $ idx + 1 + (True, False) -> performEobRun $ idx + 1 + (True, True) -> do + let newVal | coeff >= 0 = coeff + plusOne + | otherwise = coeff + minusOne + lift $ (block `MS.unsafeWrite` idx) newVal + performEobRun $ idx + 1 + + unpack idx | idx > maxIndex = pure 0 + unpack idx = do + rrrrssss <- decodeRrrrSsss $ acHuffmanTree params + case rrrrssss of + (0xF, 0) -> error "Dunno" + ( r, 0) -> do + lowBits <- unpackInt r + let newEobRun = (1 `unsafeShiftL` r) + lowBits - 1 + performEobRun idx + pure newEobRun + + ( r, _) -> do + idx' <- updateCoeffs r idx + val <- getBitVal + when (idx <= maxIndex) $ + (lift $ (block `MS.unsafeWrite` idx') val) + unpack $ idx' + 1 + + updateCoeffs r idx + | r < 0 = pure idx + | idx > maxIndex = pure idx + updateCoeffs r idx = do + coeff <- lift $ block `MS.unsafeRead` idx + if coeff /= 0 then do + bit <- getNextBitJpg + when (bit && coeff .&. plusOne == 0) $ + if coeff >= 0 then + lift . (block `MS.unsafeWrite` idx) $ coeff + plusOne + else + lift . (block `MS.unsafeWrite` idx) $ coeff + minusOne + updateCoeffs r $ idx + 1 + else + updateCoeffs (r - 1) $ idx + 1 + +type Unpacker s = + JpgUnpackerParameter -> MS.STVector s Int16 -> MutableMacroBlock s Int16 -> Int32 + -> BoolReader s Int32 + + +prepareUnpacker :: [([(JpgUnpackerParameter, a)], L.ByteString)] + -> ST s ( V.Vector (V.Vector (JpgUnpackerParameter, Unpacker s)) + , M.STVector s BoolState) +prepareUnpacker lst = do + let boolStates = V.fromList $ map snd infos + vec <- V.unsafeThaw boolStates + return (V.fromList $ map fst infos, vec) + where infos = map prepare lst + prepare ([], _) = error "progressiveUnpack, no component" + prepare (whole@((param, _) : _) , byteString) = + (V.fromList $ map (\(p,_) -> (p, unpacker)) whole, boolReader) + where unpacker = selection (successiveApprox param) (coefficientRange param) + boolReader = initBoolStateJpg . B.concat $ L.toChunks byteString + + selection (_, 0) (0, _) = decodeFirstDC + selection (_, 0) _ = decodeFirstAc + selection _ (0, _) = decodeRefineDc + selection _ _ = decodeRefineAc + +data ComponentData s = ComponentData + { componentIndices :: V.Vector (VS.Vector Int) + , componentBlocks :: V.Vector (MutableMacroBlock s Int16) + , componentId :: !Int + , componentBlockCount :: !Int + } + +-- | Iteration from to n in monadic context, without data +-- keeping. +lineMap :: (Monad m) => Int -> (Int -> m ()) -> m () +{-# INLINE lineMap #-} +lineMap count f = go 0 + where go n | n >= count = return () + go n = f n >> go (n + 1) + +progressiveUnpack :: (Int, Int) + -> JpgFrameHeader + -> V.Vector (MacroBlock Int16) + -> [([(JpgUnpackerParameter, a)], L.ByteString)] + -> ST s (MutableImage s PixelYCbCr8) +progressiveUnpack (maxiW, maxiH) frame quants lst = do + (unpackers, readers) <- prepareUnpacker lst + allBlocks <- mapM allocateWorkingBlocks . zip [0..] $ jpgComponents frame + :: ST s [ComponentData s] + let scanCount = length lst + restartIntervalValue = case lst of + ((p,_):_,_): _ -> restartInterval p + _ -> -1 + dcCoeffs <- MS.replicate imgComponentCount 0 + eobRuns <- MS.replicate (length lst) 0 + workBlock <- createEmptyMutableMacroBlock + writeIndices <- MS.replicate imgComponentCount (0 :: Int) + restartIntervals <- MS.replicate scanCount restartIntervalValue + let elementCount = imgWidth * imgHeight * fromIntegral imgComponentCount + img <- MutableImage imgWidth imgHeight <$> MS.replicate elementCount 128 + + let processRestartInterval = + forM_ [0 .. scanCount - 1] $ \ix -> do + v <- restartIntervals `MS.read` ix + if v == 0 then do + -- reset DC prediction + when (ix == 0) (MS.set dcCoeffs 0) + reader <- readers `M.read` ix + (_, updated) <- runBoolReaderWith reader $ + byteAlignJpg >> decodeRestartInterval + (readers `M.write` ix) updated + (eobRuns `MS.unsafeWrite` ix) 0 + (restartIntervals `MS.unsafeWrite` ix) $ restartIntervalValue - 1 + else + (restartIntervals `MS.unsafeWrite` ix) $ v - 1 + + + lineMap imageMcuHeight $ \mmY -> do + -- Reset all blocks to 0 + forM_ allBlocks $ V.mapM_ (flip MS.set 0) . componentBlocks + MS.set writeIndices 0 + + lineMap imageMcuWidth $ \_mmx -> do + processRestartInterval + V.forM_ unpackers $ V.mapM_ $ \(unpackParam, unpacker) -> do + boolState <- readers `M.read` readerIndex unpackParam + eobrun <- eobRuns `MS.read` readerIndex unpackParam + let componentNumber = componentIndex unpackParam + writeIndex <- writeIndices `MS.read` componentNumber + let componentData = allBlocks !! componentNumber + indexVector = + componentIndices componentData ! indiceVector unpackParam + realIndex = indexVector VS.! (writeIndex + blockIndex unpackParam) + writeBlock = + componentBlocks componentData ! realIndex + (eobrun', state) <- + runBoolReaderWith boolState $ + unpacker unpackParam dcCoeffs writeBlock eobrun + + (readers `M.write` readerIndex unpackParam) state + (eobRuns `MS.write` readerIndex unpackParam) eobrun' + + -- Update the write indices + forM_ allBlocks $ \comp -> do + writeIndex <- writeIndices `MS.read` componentId comp + let newIndex = writeIndex + componentBlockCount comp + (writeIndices `MS.write` componentId comp) $ newIndex + + forM_ allBlocks $ \compData -> do + let compBlocks = componentBlocks compData + cId = componentId compData + table = quants ! min 1 cId + comp = jpgComponents frame !! cId + compW = fromIntegral $ horizontalSamplingFactor comp + compH = fromIntegral $ verticalSamplingFactor comp + cw8 = maxiW - fromIntegral (horizontalSamplingFactor comp) + 1 + ch8 = maxiH - fromIntegral (verticalSamplingFactor comp) + 1 + + rasterMap (imageMcuWidth * compW) compH $ \rx y -> do + let ry = mmY * maxiH + y + block = compBlocks ! (y * imageMcuWidth * compW + rx) + transformed <- decodeMacroBlock table workBlock block + unpackMacroBlock imgComponentCount + cw8 ch8 cId (rx * cw8) ry + img transformed + + return img + + where imgComponentCount = length $ jpgComponents frame + mcuBlockWidth = maxiW * dctBlockSize + mcuBlockHeight = maxiH * dctBlockSize + + imgWidth = fromIntegral $ jpgWidth frame + imgHeight = fromIntegral $ jpgHeight frame + imageMcuWidth = (imgWidth + mcuBlockWidth - 1) `div` mcuBlockWidth + imageMcuHeight = (imgHeight + mcuBlockHeight - 1) `div` mcuBlockHeight + + allocateWorkingBlocks (ix, comp) = do + let blockCount = hSample * vSample * imageMcuWidth + blocks <- V.replicateM blockCount createEmptyMutableMacroBlock + return $ ComponentData + { componentBlocks = blocks + , componentIndices = createMcuLineIndices comp imageMcuWidth + , componentBlockCount = hSample * vSample + , componentId = ix + } + where hSample = fromIntegral $ horizontalSamplingFactor comp + vSample = fromIntegral $ verticalSamplingFactor comp + diff --git a/Codec/Picture/Jpg/Types.hs b/Codec/Picture/Jpg/Types.hs --- a/Codec/Picture/Jpg/Types.hs +++ b/Codec/Picture/Jpg/Types.hs @@ -1,17 +1,523 @@ +{-# LANGUAGE ScopedTypeVariables #-} module Codec.Picture.Jpg.Types( MutableMacroBlock , createEmptyMutableMacroBlock , printMacroBlock + , printPureMacroBlock + , DcCoefficient + , JpgImage( .. ) + , JpgComponent( .. ) + , JpgFrameHeader( .. ) + , JpgFrame( .. ) + , JpgFrameKind( .. ) + , JpgScanHeader( .. ) + , JpgQuantTableSpec( .. ) + , JpgHuffmanTableSpec( .. ) + , JpgImageKind( .. ) + , JpgScanSpecification( .. ) + , calculateSize + , dctBlockSize ) where +import Control.Applicative( (<$>), (<*>)) +import Control.Monad( when, replicateM, forM, forM_, unless ) import Control.Monad.ST( ST ) +import Data.Bits( (.|.), (.&.), unsafeShiftL, unsafeShiftR ) import Foreign.Storable ( Storable ) +import Data.Vector.Unboxed( (!) ) +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Storable as VS import qualified Data.Vector.Storable.Mutable as M +import qualified Data.ByteString as B +import qualified Data.ByteString.Lazy as L +import Data.Int( Int16 ) +import Data.Word(Word8, Word16 ) +import Data.Binary( Binary(..) ) + +import Data.Binary.Get( Get + , getWord8 + , getWord16be + , getByteString + , skip + , bytesRead + ) + +import Data.Binary.Put( Put + , putWord8 + , putWord16be + , putLazyByteString + ) + +import Codec.Picture.InternalHelper +import Codec.Picture.Jpg.DefaultTable + +{-import Debug.Trace-} import Text.Printf +-- | Type only used to make clear what kind of integer we are carrying +-- Might be transformed into newtype in the future +type DcCoefficient = Int16 + -- | Macroblock that can be transformed. type MutableMacroBlock s a = M.STVector s a +data JpgFrameKind = + JpgBaselineDCTHuffman + | JpgExtendedSequentialDCTHuffman + | JpgProgressiveDCTHuffman + | JpgLosslessHuffman + | JpgDifferentialSequentialDCTHuffman + | JpgDifferentialProgressiveDCTHuffman + | JpgDifferentialLosslessHuffman + | JpgExtendedSequentialArithmetic + | JpgProgressiveDCTArithmetic + | JpgLosslessArithmetic + | JpgDifferentialSequentialDCTArithmetic + | JpgDifferentialProgressiveDCTArithmetic + | JpgDifferentialLosslessArithmetic + | JpgQuantizationTable + | JpgHuffmanTableMarker + | JpgStartOfScan + | JpgEndOfImage + | JpgAppSegment Word8 + | JpgExtensionSegment Word8 + + | JpgRestartInterval + | JpgRestartIntervalEnd Word8 + deriving (Eq, Show) + +data JpgFrame = + JpgAppFrame !Word8 B.ByteString + | JpgExtension !Word8 B.ByteString + | JpgQuantTable ![JpgQuantTableSpec] + | JpgHuffmanTable ![(JpgHuffmanTableSpec, HuffmanPackedTree)] + | JpgScanBlob !JpgScanHeader !L.ByteString + | JpgScans !JpgFrameKind !JpgFrameHeader + | JpgIntervalRestart !Word16 + deriving Show + +data JpgFrameHeader = JpgFrameHeader + { jpgFrameHeaderLength :: !Word16 + , jpgSamplePrecision :: !Word8 + , jpgHeight :: !Word16 + , jpgWidth :: !Word16 + , jpgImageComponentCount :: !Word8 + , jpgComponents :: ![JpgComponent] + } + deriving Show + + +instance SizeCalculable JpgFrameHeader where + calculateSize hdr = 2 + 1 + 2 + 2 + 1 + + sum [calculateSize c | c <- jpgComponents hdr] + +data JpgComponent = JpgComponent + { componentIdentifier :: !Word8 + -- | Stored with 4 bits + , horizontalSamplingFactor :: !Word8 + -- | Stored with 4 bits + , verticalSamplingFactor :: !Word8 + , quantizationTableDest :: !Word8 + } + deriving Show + +instance SizeCalculable JpgComponent where + calculateSize _ = 3 + +data JpgImage = JpgImage { jpgFrame :: [JpgFrame] } + deriving Show + +data JpgScanSpecification = JpgScanSpecification + { componentSelector :: !Word8 + -- | Encoded as 4 bits + , dcEntropyCodingTable :: !Word8 + -- | Encoded as 4 bits + , acEntropyCodingTable :: !Word8 + + } + deriving Show + +instance SizeCalculable JpgScanSpecification where + calculateSize _ = 2 + +data JpgScanHeader = JpgScanHeader + { scanLength :: !Word16 + , scanComponentCount :: !Word8 + , scans :: [JpgScanSpecification] + + -- | (begin, end) + , spectralSelection :: (Word8, Word8) + + -- | Encoded as 4 bits + , successiveApproxHigh :: !Word8 + + -- | Encoded as 4 bits + , successiveApproxLow :: !Word8 + } + deriving Show + +instance SizeCalculable JpgScanHeader where + calculateSize hdr = 2 + 1 + + sum [calculateSize c | c <- scans hdr] + + 2 + + 1 + +data JpgQuantTableSpec = JpgQuantTableSpec + { -- | Stored on 4 bits + quantPrecision :: !Word8 + + -- | Stored on 4 bits + , quantDestination :: !Word8 + + , quantTable :: MacroBlock Int16 + } + deriving Show + +class SizeCalculable a where + calculateSize :: a -> Int + +-- | Type introduced only to avoid some typeclass overlapping +-- problem +newtype TableList a = TableList [a] + +instance (SizeCalculable a, Binary a) => Binary (TableList a) where + put (TableList lst) = do + putWord16be . fromIntegral $ sum [calculateSize table | table <- lst] + 2 + mapM_ put lst + + get = TableList <$> (getWord16be >>= \s -> innerParse (fromIntegral s - 2)) + where innerParse :: Int -> Get [a] + innerParse 0 = return [] + innerParse size = do + onStart <- fromIntegral <$> bytesRead + table <- get + onEnd <- fromIntegral <$> bytesRead + (table :) <$> innerParse (size - (onEnd - onStart)) + +instance SizeCalculable JpgQuantTableSpec where + calculateSize table = + 1 + (fromIntegral (quantPrecision table) + 1) * 64 + +instance Binary JpgQuantTableSpec where + put table = do + let precision = quantPrecision table + put4BitsOfEach precision (quantDestination table) + forM_ (VS.toList $ quantTable table) $ \coeff -> + if precision == 0 then putWord8 $ fromIntegral coeff + else putWord16be $ fromIntegral coeff + + get = do + (precision, dest) <- get4BitOfEach + coeffs <- replicateM 64 $ if precision == 0 + then fromIntegral <$> getWord8 + else fromIntegral <$> getWord16be + return JpgQuantTableSpec + { quantPrecision = precision + , quantDestination = dest + , quantTable = VS.fromListN 64 coeffs + } + +data JpgHuffmanTableSpec = JpgHuffmanTableSpec + { -- | 0 : DC, 1 : AC, stored on 4 bits + huffmanTableClass :: !DctComponent + -- | Stored on 4 bits + , huffmanTableDest :: !Word8 + + , huffSizes :: !(VU.Vector Word8) + , huffCodes :: !(V.Vector (VU.Vector Word8)) + } + deriving Show + +instance SizeCalculable JpgHuffmanTableSpec where + calculateSize table = 1 + 16 + sum [fromIntegral e | e <- VU.toList $ huffSizes table] + +instance Binary JpgHuffmanTableSpec where + put table = do + let classVal = if huffmanTableClass table == DcComponent + then 0 else 1 + put4BitsOfEach classVal $ huffmanTableDest table + mapM_ put . VU.toList $ huffSizes table + forM_ [0 .. 15] $ \i -> + when (huffSizes table ! i /= 0) + (let elements = VU.toList $ huffCodes table V.! i + in mapM_ put elements) + + get = do + (huffClass, huffDest) <- get4BitOfEach + sizes <- replicateM 16 getWord8 + codes <- forM sizes $ \s -> + VU.replicateM (fromIntegral s) getWord8 + return JpgHuffmanTableSpec + { huffmanTableClass = + if huffClass == 0 then DcComponent else AcComponent + , huffmanTableDest = huffDest + , huffSizes = VU.fromListN 16 sizes + , huffCodes = V.fromListN 16 codes + } + +instance Binary JpgImage where + put (JpgImage { jpgFrame = frames }) = + putWord8 0xFF >> putWord8 0xD8 >> mapM_ putFrame frames + >> putWord8 0xFF >> putWord8 0xD9 + + get = do + let startOfImageMarker = 0xD8 + -- endOfImageMarker = 0xD9 + checkMarker commonMarkerFirstByte startOfImageMarker + eatUntilCode + frames <- parseFrames + {-checkMarker commonMarkerFirstByte endOfImageMarker-} + return JpgImage { jpgFrame = frames } + +eatUntilCode :: Get () +eatUntilCode = do + code <- getWord8 + unless (code == 0xFF) eatUntilCode + +takeCurrentFrame :: Get B.ByteString +takeCurrentFrame = do + size <- getWord16be + getByteString (fromIntegral size - 2) + +putFrame :: JpgFrame -> Put +putFrame (JpgAppFrame appCode str) = + put (JpgAppSegment appCode) >> putWord16be (fromIntegral $ B.length str) >> put str +putFrame (JpgExtension appCode str) = + put (JpgExtensionSegment appCode) >> putWord16be (fromIntegral $ B.length str) >> put str +putFrame (JpgQuantTable tables) = + put JpgQuantizationTable >> put (TableList tables) +putFrame (JpgHuffmanTable tables) = + put JpgHuffmanTableMarker >> put (TableList $ map fst tables) +putFrame (JpgIntervalRestart size) = + put JpgRestartInterval >> put (RestartInterval size) +putFrame (JpgScanBlob hdr blob) = + put JpgStartOfScan >> put hdr >> putLazyByteString blob +putFrame (JpgScans kind hdr) = + put kind >> put hdr + +-------------------------------------------------- +---- Serialization instances +-------------------------------------------------- +commonMarkerFirstByte :: Word8 +commonMarkerFirstByte = 0xFF + +checkMarker :: Word8 -> Word8 -> Get () +checkMarker b1 b2 = do + rb1 <- getWord8 + rb2 <- getWord8 + when (rb1 /= b1 || rb2 /= b2) + (fail "Invalid marker used") + +extractScanContent :: L.ByteString -> (L.ByteString, L.ByteString) +extractScanContent str = aux 0 + where maxi = fromIntegral $ L.length str - 1 + + aux n | n >= maxi = (str, L.empty) + | v == 0xFF && vNext /= 0 && not isReset = L.splitAt n str + | otherwise = aux (n + 1) + where v = str `L.index` n + vNext = str `L.index` (n + 1) + isReset = 0xD0 <= vNext && vNext <= 0xD7 + +parseFrames :: Get [JpgFrame] +parseFrames = do + kind <- get + let parseNextFrame = do + word <- getWord8 + when (word /= 0xFF) $ do + readedData <- bytesRead + fail $ "Invalid Frame marker (" ++ show word + ++ ", bytes read : " ++ show readedData ++ ")" + parseFrames + + case kind of + JpgEndOfImage -> return [] + JpgAppSegment c -> + (\frm lst -> JpgAppFrame c frm : lst) <$> takeCurrentFrame <*> parseNextFrame + JpgExtensionSegment c -> + (\frm lst -> JpgExtension c frm : lst) <$> takeCurrentFrame <*> parseNextFrame + JpgQuantizationTable -> + (\(TableList quants) lst -> JpgQuantTable quants : lst) <$> get <*> parseNextFrame + JpgRestartInterval -> + (\(RestartInterval i) lst -> JpgIntervalRestart i : lst) <$> get <*> parseNextFrame + JpgHuffmanTableMarker -> + (\(TableList huffTables) lst -> + JpgHuffmanTable [(t, packHuffmanTree . buildPackedHuffmanTree $ huffCodes t) | t <- huffTables] : lst) + <$> get <*> parseNextFrame + JpgStartOfScan -> + (\frm imgData -> + let (d, other) = extractScanContent imgData + in + case runGet parseFrames (L.drop 1 other) of + Left _ -> [JpgScanBlob frm d] + Right lst -> JpgScanBlob frm d : lst + ) <$> get <*> getRemainingLazyBytes + + _ -> (\hdr lst -> JpgScans kind hdr : lst) <$> get <*> parseNextFrame + +buildPackedHuffmanTree :: V.Vector (VU.Vector Word8) -> HuffmanTree +buildPackedHuffmanTree = buildHuffmanTree . map VU.toList . V.toList + +secondStartOfFrameByteOfKind :: JpgFrameKind -> Word8 +secondStartOfFrameByteOfKind = aux + where + aux JpgBaselineDCTHuffman = 0xC0 + aux JpgExtendedSequentialDCTHuffman = 0xC1 + aux JpgProgressiveDCTHuffman = 0xC2 + aux JpgLosslessHuffman = 0xC3 + aux JpgDifferentialSequentialDCTHuffman = 0xC5 + aux JpgDifferentialProgressiveDCTHuffman = 0xC6 + aux JpgDifferentialLosslessHuffman = 0xC7 + aux JpgExtendedSequentialArithmetic = 0xC9 + aux JpgProgressiveDCTArithmetic = 0xCA + aux JpgLosslessArithmetic = 0xCB + aux JpgHuffmanTableMarker = 0xC4 + aux JpgDifferentialSequentialDCTArithmetic = 0xCD + aux JpgDifferentialProgressiveDCTArithmetic = 0xCE + aux JpgDifferentialLosslessArithmetic = 0xCF + aux JpgEndOfImage = 0xD9 + aux JpgQuantizationTable = 0xDB + aux JpgStartOfScan = 0xDA + aux JpgRestartInterval = 0xDD + aux (JpgRestartIntervalEnd v) = v + aux (JpgAppSegment a) = a + aux (JpgExtensionSegment a) = a + +data JpgImageKind = BaseLineDCT | ProgressiveDCT + +instance Binary JpgFrameKind where + put v = putWord8 0xFF >> put (secondStartOfFrameByteOfKind v) + get = do + -- no lookahead :( + {-word <- getWord8-} + word2 <- getWord8 + return $ case word2 of + 0xC0 -> JpgBaselineDCTHuffman + 0xC1 -> JpgExtendedSequentialDCTHuffman + 0xC2 -> JpgProgressiveDCTHuffman + 0xC3 -> JpgLosslessHuffman + 0xC4 -> JpgHuffmanTableMarker + 0xC5 -> JpgDifferentialSequentialDCTHuffman + 0xC6 -> JpgDifferentialProgressiveDCTHuffman + 0xC7 -> JpgDifferentialLosslessHuffman + 0xC9 -> JpgExtendedSequentialArithmetic + 0xCA -> JpgProgressiveDCTArithmetic + 0xCB -> JpgLosslessArithmetic + 0xCD -> JpgDifferentialSequentialDCTArithmetic + 0xCE -> JpgDifferentialProgressiveDCTArithmetic + 0xCF -> JpgDifferentialLosslessArithmetic + 0xD9 -> JpgEndOfImage + 0xDA -> JpgStartOfScan + 0xDB -> JpgQuantizationTable + 0xDD -> JpgRestartInterval + a | a >= 0xF0 -> JpgExtensionSegment a + | a >= 0xE0 -> JpgAppSegment a + | a >= 0xD0 && a <= 0xD7 -> JpgRestartIntervalEnd a + | otherwise -> error ("Invalid frame marker (" ++ show a ++ ")") + +put4BitsOfEach :: Word8 -> Word8 -> Put +put4BitsOfEach a b = put $ (a `unsafeShiftL` 4) .|. b + +get4BitOfEach :: Get (Word8, Word8) +get4BitOfEach = do + val <- get + return ((val `unsafeShiftR` 4) .&. 0xF, val .&. 0xF) + +newtype RestartInterval = RestartInterval Word16 + +instance Binary RestartInterval where + put (RestartInterval i) = putWord16be 4 >> putWord16be i + get = do + size <- getWord16be + when (size /= 4) (fail "Invalid jpeg restart interval size") + RestartInterval <$> getWord16be + +instance Binary JpgComponent where + get = do + ident <- getWord8 + (horiz, vert) <- get4BitOfEach + quantTableIndex <- getWord8 + return JpgComponent + { componentIdentifier = ident + , horizontalSamplingFactor = horiz + , verticalSamplingFactor = vert + , quantizationTableDest = quantTableIndex + } + put v = do + put $ componentIdentifier v + put4BitsOfEach (horizontalSamplingFactor v) $ verticalSamplingFactor v + put $ quantizationTableDest v + +instance Binary JpgFrameHeader where + get = do + beginOffset <- fromIntegral <$> bytesRead + frmHLength <- getWord16be + samplePrec <- getWord8 + h <- getWord16be + w <- getWord16be + compCount <- getWord8 + components <- replicateM (fromIntegral compCount) get + endOffset <- fromIntegral <$> bytesRead + when (beginOffset - endOffset < fromIntegral frmHLength) + (skip $ fromIntegral frmHLength - (endOffset - beginOffset)) + return JpgFrameHeader + { jpgFrameHeaderLength = frmHLength + , jpgSamplePrecision = samplePrec + , jpgHeight = h + , jpgWidth = w + , jpgImageComponentCount = compCount + , jpgComponents = components + } + + put v = do + putWord16be $ jpgFrameHeaderLength v + putWord8 $ jpgSamplePrecision v + putWord16be $ jpgHeight v + putWord16be $ jpgWidth v + putWord8 $ jpgImageComponentCount v + mapM_ put $ jpgComponents v + +instance Binary JpgScanSpecification where + put v = do + put $ componentSelector v + put4BitsOfEach (dcEntropyCodingTable v) $ acEntropyCodingTable v + + get = do + compSel <- get + (dc, ac) <- get4BitOfEach + return JpgScanSpecification { + componentSelector = compSel + , dcEntropyCodingTable = dc + , acEntropyCodingTable = ac + } + +instance Binary JpgScanHeader where + get = do + thisScanLength <- getWord16be + compCount <- getWord8 + comp <- replicateM (fromIntegral compCount) get + specBeg <- get + specEnd <- get + (approxHigh, approxLow) <- get4BitOfEach + + return JpgScanHeader { + scanLength = thisScanLength, + scanComponentCount = compCount, + scans = comp, + spectralSelection = (specBeg, specEnd), + successiveApproxHigh = approxHigh, + successiveApproxLow = approxLow + } + + put v = do + putWord16be $ scanLength v + putWord8 $ scanComponentCount v + mapM_ put $ scans v + putWord8 . fst $ spectralSelection v + putWord8 . snd $ spectralSelection v + put4BitsOfEach (successiveApproxHigh v) $ successiveApproxLow v + {-# INLINE createEmptyMutableMacroBlock #-} -- | Create a new macroblock with the good array size createEmptyMutableMacroBlock :: (Storable a, Num a) => ST s (MutableMacroBlock s a) @@ -26,3 +532,15 @@ vn <- pLn (i+1) return $ printf (if i `mod` 8 == 0 then "\n%5d " else "%5d ") v ++ vn +printPureMacroBlock :: (Storable a, PrintfArg a) => MacroBlock a -> String +printPureMacroBlock block = pLn 0 + where pLn 64 = "===============================\n" + pLn i = str ++ pLn (i + 1) + where str | i `mod` 8 == 0 = printf "\n%5d " v + | otherwise = printf "%5d" v + v = block VS.! i + + +{-# INLINE dctBlockSize #-} +dctBlockSize :: Num a => a +dctBlockSize = 8 diff --git a/Codec/Picture/Tiff.hs b/Codec/Picture/Tiff.hs --- a/Codec/Picture/Tiff.hs +++ b/Codec/Picture/Tiff.hs @@ -334,7 +334,7 @@ instance BinaryParam Endianness TiffTag where getP endianness = tagOfWord16 <$> getP endianness - putP endianness = putP endianness . word16OfTag + putP endianness = putP endianness . word16OfTag data ExtendedDirectoryData = ExtendedDataNone @@ -406,11 +406,11 @@ unLong _ (ExtendedDataLong v) = pure v unLong errMessage _ = fail errMessage -cleanImageFileDirectory :: ImageFileDirectory -> ImageFileDirectory -cleanImageFileDirectory ifd@(ImageFileDirectory { ifdCount = 1 }) = aux $ ifdType ifd +cleanImageFileDirectory :: Endianness -> ImageFileDirectory -> ImageFileDirectory +cleanImageFileDirectory EndianBig ifd@(ImageFileDirectory { ifdCount = 1 }) = aux $ ifdType ifd where aux TypeShort = ifd { ifdOffset = ifdOffset ifd `unsafeShiftR` 16 } aux _ = ifd -cleanImageFileDirectory ifd = ifd +cleanImageFileDirectory _ ifd = ifd instance BinaryParam Endianness ImageFileDirectory where getP endianness = @@ -433,7 +433,7 @@ rez <- replicateM (fromIntegral count) $ getP endianness _ <- getP endianness :: Get Word32 pure rez - + putP endianness lst = do let count = fromIntegral $ length lst :: Word16 putP endianness count @@ -490,8 +490,10 @@ findIFDExtDefaultData d tag lst = case [v | v <- lst, ifdIdentifier v == tag] of [] -> pure d - (x:_) -> V.toList <$> - unLong ("Can't unlong " ++ (show tag)) (ifdExtended x) + (ImageFileDirectory { ifdExtended = ExtendedDataNone }:_) -> return d + (x:_) -> V.toList <$> unLong errorMessage (ifdExtended x) + where errorMessage = + "Can't parse tag " ++ show tag ++ " " ++ show (ifdExtended x) -- It's temporary, remove once tiff decoding is better -- handled. @@ -688,6 +690,45 @@ looperBe (writeIndex + stride) (readIndex + 2) +instance Unpackable Word32 where + type StorageType Word32 = Word32 + + offsetStride _ _ _ = (0, 1) + outAlloc _ = M.new + allocTempBuffer _ _ s = M.new $ s * 4 + mergeBackTempBuffer _ EndianLittle tempVec _ index size stride outVec = + looperLe index 0 + where looperLe _ readIndex | readIndex >= fromIntegral size = pure () + looperLe writeIndex readIndex = do + v1 <- tempVec `M.read` readIndex + v2 <- tempVec `M.read` (readIndex + 1) + v3 <- tempVec `M.read` (readIndex + 2) + v4 <- tempVec `M.read` (readIndex + 3) + let finalValue = + (fromIntegral v4 `unsafeShiftL` 24) .|. + (fromIntegral v3 `unsafeShiftL` 16) .|. + (fromIntegral v2 `unsafeShiftL` 8) .|. + fromIntegral v1 + (outVec `M.write` writeIndex) finalValue + + looperLe (writeIndex + stride) (readIndex + 4) + mergeBackTempBuffer _ EndianBig tempVec _ index size stride outVec = + looperBe index 0 + where looperBe _ readIndex | readIndex >= fromIntegral size = pure () + looperBe writeIndex readIndex = do + v1 <- tempVec `M.read` readIndex + v2 <- tempVec `M.read` (readIndex + 1) + v3 <- tempVec `M.read` (readIndex + 2) + v4 <- tempVec `M.read` (readIndex + 3) + let finalValue = + (fromIntegral v1 `unsafeShiftL` 24) .|. + (fromIntegral v2 `unsafeShiftL` 16) .|. + (fromIntegral v3 `unsafeShiftL` 8) .|. + fromIntegral v4 + (outVec `M.write` writeIndex) finalValue + + looperBe (writeIndex + stride) (readIndex + 4) + data Pack4 = Pack4 instance Unpackable Pack4 where @@ -789,7 +830,7 @@ inner (readIdx + 3) (writeIdx + 2 * stride) (line - 2) -data YCbCrSubsampling = YCbCrSubsampling +data YCbCrSubsampling = YCbCrSubsampling { ycbcrWidth :: !Int , ycbcrHeight :: !Int , ycbcrImageWidth :: !Int @@ -832,7 +873,7 @@ return $ readIndex + 1 foldM_ writer readIdx pixelIndices - + return $ readIdx + blockSize gatherStrips :: ( Unpackable comp @@ -951,8 +992,8 @@ ifdSize = 12 ifdCountSize = 2 nextOffsetSize = 4 - startExtended = initialOffset - + ifdElementCount * ifdSize + startExtended = initialOffset + + ifdElementCount * ifdSize + ifdCountSize + nextOffsetSize updater ix ifd@(ImageFileDirectory { ifdExtended = ExtendedDataAscii b }) = @@ -980,10 +1021,10 @@ ifdShorts TagBitsPerSample $ tiffBitsPerSample nfo ifdSingleLong TagSamplesPerPixel $ tiffSampleCount nfo ifdSingleLong TagRowPerStrip $ tiffRowPerStrip nfo - ifdShort TagPhotometricInterpretation - . packPhotometricInterpretation + ifdShort TagPhotometricInterpretation + . packPhotometricInterpretation $ tiffColorspace nfo - ifdShort TagPlanarConfiguration + ifdShort TagPlanarConfiguration . constantToPlaneConfiguration $ tiffPlaneConfiguration nfo ifdShort TagCompression . packCompression $ tiffCompression nfo @@ -1005,7 +1046,7 @@ skip . fromIntegral $ fromIntegral (hdrOffset hdr) - readed let endian = hdrEndianness hdr - ifd <- fmap cleanImageFileDirectory <$> getP endian + ifd <- fmap (cleanImageFileDirectory endian) <$> getP endian cleaned <- fetchExtended endian ifd let dataFind str tag = findIFDData str tag cleaned @@ -1036,13 +1077,15 @@ <*> (V.fromList <$> extDefault [2, 2] TagYCbCrSubsampling) unpack :: B.ByteString -> TiffInfo -> Either String DynamicImage +-- | while mandatory some images don't put correct +-- rowperstrip. So replacing 0 with actual image height. +unpack file nfo@TiffInfo { tiffRowPerStrip = 0 } = + unpack file $ nfo { tiffRowPerStrip = tiffHeight nfo } unpack file nfo@TiffInfo { tiffColorspace = TiffPaleted , tiffBitsPerSample = lst , tiffSampleFormat = format , tiffPalette = Just p - , tiffRowPerStrip = rowPerStrip } - | rowPerStrip == 0 = fail "Invalid row per strip" | lst == V.singleton 8 && format == [TiffSampleUint] = let applyPalette = pixelMap (\v -> pixelAt p (fromIntegral v) 0) gathered :: Image Pixel8 @@ -1065,20 +1108,15 @@ pure . ImageRGB16 $ applyPalette gathered unpack file nfo@TiffInfo { tiffColorspace = TiffCMYK - , tiffRowPerStrip = rowPerStrip , tiffBitsPerSample = lst , tiffSampleFormat = format } - | rowPerStrip == 0 = fail "Invalid row per strip" | lst == V.fromList [8, 8, 8, 8] && all (TiffSampleUint ==) format = pure . ImageCMYK8 $ gatherStrips (0 :: Word8) file nfo | lst == V.fromList [16, 16, 16, 16] && all (TiffSampleUint ==) format = pure . ImageCMYK16 $ gatherStrips (0 :: Word16) file nfo -unpack file nfo@TiffInfo { tiffColorspace = TiffMonochromeWhite0 - , tiffRowPerStrip = rowPerStrip - } - | rowPerStrip == 0 = fail "Invalid row per strip" +unpack file nfo@TiffInfo { tiffColorspace = TiffMonochromeWhite0 } | otherwise = do img <- unpack file (nfo { tiffColorspace = TiffMonochrome }) case img of @@ -1087,10 +1125,8 @@ _ -> pure img unpack file nfo@TiffInfo { tiffColorspace = TiffMonochrome - , tiffRowPerStrip = rowPerStrip , tiffBitsPerSample = lst , tiffSampleFormat = format } - | rowPerStrip == 0 = fail "Invalid row per strip" | lst == V.singleton 2 && all (TiffSampleUint ==) format = pure . ImageY8 . pixelMap (colorMap (64 *)) $ gatherStrips Pack2 file nfo | lst == V.singleton 4 && all (TiffSampleUint ==) format = @@ -1101,13 +1137,15 @@ pure . ImageY16 . pixelMap (16 *) $ gatherStrips Pack12 file nfo | lst == V.singleton 16 && all (TiffSampleUint ==) format = pure . ImageY16 $ gatherStrips (0 :: Word16) file nfo + | lst == V.singleton 32 && all (TiffSampleUint ==) format = + pure . ImageY16 $ pixelMap (toWord16) img + where toWord16 v = fromIntegral $ v `unsafeShiftR` 16 + img = gatherStrips (0 :: Word32) file nfo :: Image Pixel32 unpack file nfo@TiffInfo { tiffColorspace = TiffYCbCr , tiffBitsPerSample = lst - , tiffRowPerStrip = rowPerStrip , tiffPlaneConfiguration = PlanarConfigContig , tiffSampleFormat = format } - | rowPerStrip == 0 = fail "Invalid row per strip" | lst == V.fromList [8, 8, 8] && all (TiffSampleUint ==) format = pure . ImageYCbCr8 $ gatherStrips cbcrConf file nfo where defaulting 0 = 2 @@ -1123,10 +1161,8 @@ } unpack file nfo@TiffInfo { tiffColorspace = TiffRGB - , tiffRowPerStrip = rowPerStrip , tiffBitsPerSample = lst , tiffSampleFormat = format } - | rowPerStrip == 0 = fail "Invalid row per strip" | lst == V.fromList [2, 2, 2] && all (TiffSampleUint ==) format = pure . ImageRGB8 . pixelMap (colorMap (64 *)) $ gatherStrips Pack2 file nfo | lst == V.fromList [4, 4, 4] && all (TiffSampleUint ==) format = @@ -1140,10 +1176,8 @@ | lst == V.fromList [16, 16, 16, 16] && all (TiffSampleUint ==) format = pure . ImageRGBA16 $ gatherStrips (0 :: Word16) file nfo unpack file nfo@TiffInfo { tiffColorspace = TiffMonochrome - , tiffRowPerStrip = rowPerStrip , tiffBitsPerSample = lst , tiffSampleFormat = format } - | rowPerStrip == 0 = fail "Invalid row per strip" -- some files are a little bit borked... | lst == V.fromList [8, 8, 8] && all (TiffSampleUint ==) format = pure . ImageRGB8 $ gatherStrips (0 :: Word8) file nfo @@ -1211,7 +1245,7 @@ encodeTiff :: forall px. (TiffSaveable px) => Image px -> Lb.ByteString encodeTiff img = runPut $ putP rawPixelData hdr where intSampleCount = componentCount (undefined :: px) - sampleCount = fromIntegral intSampleCount + sampleCount = fromIntegral intSampleCount sampleType = (undefined :: PixelBaseComponent px) pixelData = imageData img @@ -1226,7 +1260,7 @@ headerSize = 8 hdr = TiffInfo - { tiffHeader = TiffHeader + { tiffHeader = TiffHeader { hdrEndianness = EndianLittle , hdrOffset = headerSize + imageSize } diff --git a/Codec/Picture/Types.hs b/Codec/Picture/Types.hs --- a/Codec/Picture/Types.hs +++ b/Codec/Picture/Types.hs @@ -22,6 +22,7 @@ -- ** Pixel types , Pixel8 , Pixel16 + , Pixel32 , PixelF , PixelYA8( .. ) , PixelYA16( .. ) @@ -79,7 +80,7 @@ import Control.Monad.Primitive ( PrimMonad, PrimState ) import Foreign.Storable ( Storable ) import Data.Bits( unsafeShiftL, unsafeShiftR ) -import Data.Word( Word8, Word16 ) +import Data.Word( Word8, Word16, Word32 ) import Data.List( foldl' ) import Data.Vector.Storable ( (!) ) import qualified Data.Vector.Storable as V @@ -361,6 +362,9 @@ -- | Simple alias for greyscale value in 16 bits. type Pixel16 = Word16 +-- | Simple alias for greyscale value in 16 bits. +type Pixel32 = Word32 + -- | Floating greyscale value, the 0 to 255 8 bit range maps -- to 0 to 1 in this floating version type PixelF = Float @@ -767,6 +771,11 @@ computeLuma = id extractLumaPlane = id +instance LumaPlaneExtractable Pixel32 where + {-# INLINE computeLuma #-} + computeLuma = id + extractLumaPlane = id + instance LumaPlaneExtractable PixelF where {-# INLINE computeLuma #-} computeLuma = id @@ -885,6 +894,31 @@ instance ColorConvertible Pixel16 PixelRGBA16 where {-# INLINE promotePixel #-} promotePixel c = PixelRGBA16 c c c maxBound + +-------------------------------------------------- +---- Pixel32 instances +-------------------------------------------------- +instance Pixel Pixel32 where + type PixelBaseComponent Pixel32 = Word32 + + {-# INLINE mixWith #-} + mixWith f = f 0 + + {-# INLINE colorMap #-} + colorMap f = f + + componentCount _ = 1 + pixelAt (Image { imageWidth = w, imageData = arr }) x y = arr ! (x + y * w) + + readPixel image@(MutableImage { mutableImageData = arr }) x y = + arr `M.read` mutablePixelBaseIndex image x y + + writePixel image@(MutableImage { mutableImageData = arr }) x y = + arr `M.write` mutablePixelBaseIndex image x y + + unsafePixelAt = V.unsafeIndex + unsafeReadPixel = M.unsafeRead + unsafeWritePixel = M.unsafeWrite -------------------------------------------------- ---- PixelF instances diff --git a/JuicyPixels.cabal b/JuicyPixels.cabal --- a/JuicyPixels.cabal +++ b/JuicyPixels.cabal @@ -1,5 +1,5 @@ Name: JuicyPixels -Version: 3.1.1.1 +Version: 3.1.2 Synopsis: Picture loading/serialization (in png, jpeg, bitmap, gif, tiff and radiance) Description: <> @@ -15,8 +15,10 @@ Build-type: Simple -- Constraint on the version of Cabal needed to build this package. -Cabal-version: >= 1.10 +Cabal-version: >= 1.16 +Extra-source-files: changelog + Source-Repository head Type: git Location: git://github.com/Twinside/Juicy.Pixels.git @@ -24,7 +26,7 @@ Source-Repository this Type: git Location: git://github.com/Twinside/Juicy.Pixels.git - Tag: v3.1.1.1 + Tag: v3.1.2 Flag Mmap Description: Enable the file loading via mmap (memory map) @@ -46,7 +48,7 @@ Build-depends: base >= 4 && < 5, bytestring >= 0.9 && < 0.11, mtl >= 1.1 && < 2.2, - binary >= 0.6.4.0 && < 0.7.2, + binary >= 0.6.4.0 && < 0.8, zlib >= 0.5.3.1 && < 0.6, transformers >= 0.2.2 && < 0.4, vector >= 0.9 && < 0.11, @@ -62,6 +64,8 @@ Codec.Picture.Jpg.FastIdct, Codec.Picture.Jpg.FastDct, Codec.Picture.Jpg.Types, + Codec.Picture.Jpg.Common, + Codec.Picture.Jpg.Progressive, Codec.Picture.Gif.LZW, Codec.Picture.Png.Export, Codec.Picture.Png.Type, diff --git a/changelog b/changelog new file mode 100644 --- /dev/null +++ b/changelog @@ -0,0 +1,71 @@ +-*-change-log-*- + +v3.1.2 December 2013 + * Adding support for progressive jpeg. + * Adding support for plane separated MCU jpeg. + * Adding support for grayscale 32bit decoding (with reduced precision to + 16bits). + * Fixing erroneous bitmap decoding in case of excessive padding (#31). + +v3.1.1.1 October 2013 + * Fixing some spurious crash while decoding some JPEG image (#30). + +v3.1.1 October 2013 + * Adding uncompresed TIFF saving. + * Adding error message to avoid loading progressive loading. + * Made MMAP optional, controled by a cabal flag. + * adding dynamicPixelMap helper function. + * Handling png transparency using color key (#26). + +v3.1 June 2013 + * Adding basic handling of 16bits pixel types. + * Addition of Tiff reading: + - 2, 4, 8, 16 bit depth reading (planar and contiguous for each). + - CMYK, YCbCr, RGB, Paletted, Greyscale. + - Uncompressed, PackBits, LZW. + * Some new tiny helper functions (nothing too fancy). + * Huge performances improvement. + +v3.0 January 2013 + * Simplification of the 'Pixel' typeclass, removed many unused part. + * Removal of some Storable instances for pixel types. + * Amelioration of the documentation. + * Support for High Dynamic range images, opening support for different pixel + base component. + * Support for the Radiance file format (or RGBE, file extension .pic and .hdr). + * Dropped the cereal library in favor of the last version of Binary, present + in the Haskell platform. Every dependencies are now present in the platform. + +v2.0.2 October 2012 + * Decoding of interleaved gif image. + * Decoding delta coded gif animation. + * Bumping dependencies. + +v2.0.1 September 2012 + * Documentation enhancements. + * Fixing some huge gif file loading. + * Fixing performance problem of Bitmap and Jpeg savings. + +v2.0 September 2012 + * New extractComponent version with type safe plane extraction. + * Gif file reading. + +v1.3 June 2012 + * Fix extractComponent function. + * Adding saving for YA8 functions. + +v1.2.1 April 2012 + * Dependencies version bump. + * Dropping array dependency. + +v1.2 March 2012 + * Adding a generateImage helper function. + * Adding NFData instances for image. + * Adding JPEG writing. + +v1.1 February 2012 + * Switching to vector for arrays, big performance improvement. + +v1.0 January 2012 + * Initial release +