Fix #85 correctly detect colorspace for Jpeg (didn't knew there were …

…others...)

changelog

@@ -1,6 +1,7 @@
 -*-change-log-*-
 
 v3.2.3 February 2015
+ * Adding: Support for RGB in jpeg (yeah, that exist...)
  * Addinng: a pixelFoldMap functions analog to the `foldMap` function
   of the Foldable type class.
 

src/Codec/Picture/Jpg.hs

@@ -235,6 +235,7 @@ data JpgDecoderState = JpgDecoderState
     , currentRestartInterv  :: !Int
     , currentFrame          :: Maybe JpgFrameHeader
     , componentIndexMapping :: ![(Word8, Int)]
+    , outputColorSpace      :: !JpgColorSpace
     , isProgressive         :: !Bool
     , maximumHorizontalResolution :: !Int
     , maximumVerticalResolution   :: !Int
@@ -259,15 +260,33 @@ emptyDecoderState = JpgDecoderState
     , currentRestartInterv = -1
     , currentFrame         = Nothing
     , componentIndexMapping = []
+    , outputColorSpace = JpgColorSpaceYCbCr
     , isProgressive        = False
     , maximumHorizontalResolution = 0
     , maximumVerticalResolution   = 0
     , seenBlobs = 0
     }
 
+colorSpaceOfComponentStr :: String -> JpgColorSpace
+colorSpaceOfComponentStr s = case s of
+  [_] -> JpgColorSpaceY
+  [_,_] -> JpgColorSpaceYA
+  "\1\2\3" -> JpgColorSpaceYCbCr
+  "RGB" -> JpgColorSpaceRGB
+  "YCc" -> JpgColorSpaceYCC
+  [_,_,_] -> JpgColorSpaceYCbCr
+
+  "RGBA" -> JpgColorSpaceRGBA
+  "YCcA" -> JpgColorSpaceYCCA
+  "CMYK" -> JpgColorSpaceCMYK
+  "YCcK" -> JpgColorSpaceYCCK
+  [_,_,_,_] -> JpgColorSpaceCMYK
+  _ -> JpgColorSpaceYCbCr
+
 -- | This pseudo interpreter interpret the Jpg frame for the huffman,
 -- quant table and restart interval parameters.
 jpgMachineStep :: JpgFrame -> JpgScripter s ()
+jpgMachineStep (JpgAdobeAPP14 _) = pure ()
 jpgMachineStep (JpgAppFrame _ _) = pure ()
 jpgMachineStep (JpgExtension _ _) = pure ()
 jpgMachineStep (JpgScanBlob hdr raw_data) = do
@@ -339,6 +358,10 @@ jpgMachineStep (JpgScans kind hdr) = modify $ \s ->
    s { currentFrame = Just hdr
      , componentIndexMapping =
           [(componentIdentifier comp, ix) | (ix, comp) <- zip [0..] $ jpgComponents hdr]
+     , outputColorSpace =
+          colorSpaceOfComponentStr
+            [toEnum . fromEnum $ componentIdentifier comp
+                    | comp <- jpgComponents hdr]
      , isProgressive = case kind of
             JpgProgressiveDCTHuffman -> True
             _ -> False
@@ -456,6 +479,16 @@ gatherImageKind lst = case [k | JpgScans k _ <- lst, isDctSpecifier k] of
 gatherScanInfo :: JpgImage -> (JpgFrameKind, JpgFrameHeader)
 gatherScanInfo img = head [(a, b) | JpgScans a b <- jpgFrame img]
 
+dynamicOfColorSpace :: (Monad m) => Int -> Int -> (JpgColorSpace, VS.Vector Word8)
+                    -> m DynamicImage
+dynamicOfColorSpace w h (color, imgData) = case color of
+  JpgColorSpaceYCbCr -> return . ImageYCbCr8 $ Image w h imgData
+  JpgColorSpaceRGB -> return . ImageRGB8 $ Image w h imgData
+  JpgColorSpaceYA -> return . ImageYA8 $ Image w h imgData
+  JpgColorSpaceY -> return . ImageY8 $ Image w h imgData
+  colorSpace -> fail $ "Wrong color space : " ++ show colorSpace
+
+
 -- | 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'
@@ -465,49 +498,58 @@ gatherScanInfo img = head [(a, b) | JpgScans a b <- jpgFrame img]
 --
 --    * PixelY8
 --
+--    * PixelYA8
+--
+--    * PixelRGB8
+--
 --    * 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
-
-            decodeProgressive = runST $ do
-                let (st, wrotten) =
-                        execRWS (mapM_ jpgMachineStep (jpgFrame img)) () emptyDecoderState
-                    Just fHdr = currentFrame st
-                progressiveUnpack
-                    (maximumHorizontalResolution st, maximumVerticalResolution st)
-                    fHdr
-                    (quantizationMatrices st)
-                    wrotten >>= unsafeFreezeImage
-
-            pixelData = runST $ do
-                let (st, wrotten) =
-                        execRWS (mapM_ jpgMachineStep (jpgFrame img)) () emptyDecoderState
-                    Just fHdr = currentFrame st
-                resultImage <- M.new imageSize
-                let wrapped = MutableImage imgWidth imgHeight resultImage
-                decodeImage 
-                    fHdr
-                    (quantizationMatrices st)
-                    wrotten
-                    wrapped
-                VS.unsafeFreeze resultImage
+  Right img -> case imgKind of
+     Just BaseLineDCT ->
+         dynamicOfColorSpace imgWidth imgHeight decodeBaseline
+     Just ProgressiveDCT -> 
+         dynamicOfColorSpace imgWidth imgHeight decodeProgressive
+     _ -> Left "Unkown JPG kind"
+    where
+      compCount = length $ jpgComponents scanInfo
+      (_,scanInfo) = gatherScanInfo img
+
+      imgKind = gatherImageKind $ jpgFrame img
+      imgWidth = fromIntegral $ jpgWidth scanInfo
+      imgHeight = fromIntegral $ jpgHeight scanInfo
+
+      imageSize = imgWidth * imgHeight * compCount
+
+      decodeProgressive = runST $ do
+          let (st, wrotten) =
+                  execRWS (mapM_ jpgMachineStep (jpgFrame img)) () emptyDecoderState
+              Just fHdr = currentFrame st
+          fimg <-
+              progressiveUnpack
+                  (maximumHorizontalResolution st, maximumVerticalResolution st)
+                  fHdr
+                  (quantizationMatrices st)
+                  wrotten
+          frozen <- unsafeFreezeImage fimg
+          return (outputColorSpace st, imageData frozen)
+
+
+      decodeBaseline = runST $ do
+          let (st, wrotten) =
+                  execRWS (mapM_ jpgMachineStep (jpgFrame img)) () emptyDecoderState
+              Just fHdr = currentFrame st
+          resultImage <- M.new imageSize
+          let wrapped = MutableImage imgWidth imgHeight resultImage
+          decodeImage 
+              fHdr
+              (quantizationMatrices st)
+              wrotten
+              wrapped
+          frozen <- VS.unsafeFreeze resultImage
+          return (outputColorSpace st, frozen)
 
 extractBlock :: Image PixelYCbCr8       -- ^ Source image
              -> MutableMacroBlock s Int16      -- ^ Mutable block where to put extracted block

src/Codec/Picture/Jpg/Types.hs

@@ -16,6 +16,7 @@ module Codec.Picture.Jpg.Types( MutableMacroBlock
                               , JpgHuffmanTableSpec( .. )
                               , JpgImageKind( .. )
                               , JpgScanSpecification( .. )
+                              , JpgColorSpace( .. )
                               , calculateSize
                               , dctBlockSize
                               ) where
@@ -96,6 +97,7 @@ data JpgFrameKind =
 
 data JpgFrame =
       JpgAppFrame        !Word8 B.ByteString
+    | JpgAdobeAPP14      !JpgAdobeApp14
     | JpgExtension       !Word8 B.ByteString
     | JpgQuantTable      ![JpgQuantTableSpec]
     | JpgHuffmanTable    ![(JpgHuffmanTableSpec, HuffmanPackedTree)]
@@ -104,6 +106,37 @@ data JpgFrame =
     | JpgIntervalRestart !Word16
     deriving Show
 
+data JpgColorSpace
+  = JpgColorSpaceYCbCr
+  | JpgColorSpaceYCC
+  | JpgColorSpaceY
+  | JpgColorSpaceYA
+  | JpgColorSpaceYCCA
+  | JpgColorSpaceYCCK
+  | JpgColorSpaceCMYK
+  | JpgColorSpaceRGB
+  | JpgColorSpaceRGBA
+  deriving Show
+
+data JpgAdobeApp14 = JpgAdobeApp14
+  { _adobeDctVersion :: !Word8
+  , _adobeFlag0      :: !Word16
+  , _adobeFlag1      :: !Word16
+  , _adobeTransform  :: !Word8
+  }
+  deriving Show
+
+instance Binary JpgAdobeApp14 where
+  get = JpgAdobeApp14
+     <$> getWord8 <*> getWord16be <*> getWord16be <*> getWord8
+
+  put (JpgAdobeApp14 v f0 f1 t) = do
+    putWord8 v
+    putWord16be f0
+    putWord16be f1
+    putWord8 t
+
+
 data JpgFrameHeader = JpgFrameHeader
     { jpgFrameHeaderLength   :: !Word16
     , jpgSamplePrecision     :: !Word8
@@ -288,6 +321,7 @@ takeCurrentFrame = do
     getByteString (fromIntegral size - 2)
 
 putFrame :: JpgFrame -> Put
+putFrame (JpgAdobeAPP14 _adobe) = return ()
 putFrame (JpgAppFrame appCode str) =
     put (JpgAppSegment appCode) >> putWord16be (fromIntegral $ B.length str) >> put str
 putFrame (JpgExtension appCode str) =
@@ -327,6 +361,12 @@ extractScanContent str = aux 0
                   vNext = str `L.index` (n + 1)
                   isReset = 0xD0 <= vNext && vNext <= 0xD7
 
+parseAdobe14 :: B.ByteString -> [JpgFrame] -> [JpgFrame]
+parseAdobe14 str lst = go where
+  go = case runGetStrict get str of
+    Left _err -> lst
+    Right app14 -> JpgAdobeAPP14 app14 : lst
+
 parseFrames :: Get [JpgFrame]
 parseFrames = do
     kind <- get
@@ -340,6 +380,8 @@ parseFrames = do
 
     case kind of
         JpgEndOfImage -> return []
+        JpgAppSegment 14 ->
+            parseAdobe14 <$> takeCurrentFrame <*> parseNextFrame
         JpgAppSegment c ->
             (\frm lst -> JpgAppFrame c frm : lst) <$> takeCurrentFrame <*> parseNextFrame
         JpgExtensionSegment c ->

test-src/main.hs

@@ -493,7 +493,6 @@ testSuite = do
     mapM_ (imgToImg . (("tests" </> "bmp") </>)) bmpValidTests
     mapM_ (imgToImg . (("tests" </> "pngsuite") </>)) ("huge.png" : validTests)
     mapM_ (imgToImg . (("tests" </> "jpeg") </>)) ("huge.jpg" : jpegValidTests)
-    mapM_ (imgToImg . (("tests" </> "jpeg") </>)) jpegValidTests
     mapM_ (radianceToBitmap . (("tests" </> "radiance") </>)) radianceTest
     mapM_ (gifToImg . (("tests" </> "gif") </>)) gifTest
     mapM_ (imgToImg . (("tests" </> "tiff") </>)) tiffValidTests