module Graphics.JPEG where
import Data.Char(chr,ord)
import Data.Word
import Data.Int
import Data.List(transpose)
import Data.Bits(testBit)
import Control.Monad(sequence, replicateM, liftM, liftM2)
import Debug.Trace
import System.IO
import Control.Monad.State(State(..), evalState, get, put)
infixr 9 `o`
o :: (c->d) -> (a->b->c) -> (a->b->d)
(g `o` f) x y = g (f x y)
type Table a = Int -> a
subst :: Eq a => a -> b -> (a->b) -> (a->b)
subst i e t j | i==j = e
| otherwise = t j
multi :: Int -> [a] -> [a]
multi n = concatMap (replicate n)
ceilDiv :: Int -> Int -> Int
ceilDiv n d = (n+d1)`div`d
data PixelRGB = PixelRGB { red, green, blue :: !Word8 }
sane :: Int -> Int
sane x = (0 `max` x) `min` 255
yCbCr2RGB :: [Int8] -> PixelRGB
yCbCr2RGB [y,cb,cr]
= let yi :: Int
yi = fromIntegral y
cbi :: Int
cbi = fromIntegral cb
cri :: Int
cri = fromIntegral cr
r :: Word8
r = fromIntegral (sane (128 + yi + 8*cri `div` 5))
g :: Word8
g = fromIntegral (sane (128 + yi cbi `div`3 4*cri `div` 5))
b :: Word8
b = fromIntegral (sane (128 + yi + cbi+cbi))
in PixelRGB r g b
yCbCr2RGB _ = error "yCbCr2RGB needs 3 elements"
type Dim = (Int,Int)
type Mat a = [[a]]
matapply :: Num a => Mat a -> [a] -> [a]
matapply m v = map (inprod v) m
inprod :: Num a => [a] -> [a] -> a
inprod = sum `o` zipWith (*)
matmap :: (a->b) -> Mat a -> Mat b
matmap = map . map
matconcat :: Mat (Mat a) -> Mat a
matconcat = concatMap (map concat . transpose)
matzip :: [Mat a] -> Mat [a]
matzip = map transpose . transpose
type Bits = [Bool]
byte2bits :: Int -> Bits
byte2bits x = map (testBit x) [7,6..0]
string2bits :: String -> Bits
string2bits = concatMap (byte2bits . ord)
byte2nibs :: Int -> (Int,Int)
byte2nibs x = x `divMod` 16
data Tree a = Nil
| Tip a
| Bin (Tree a) (Tree a)
instance Functor Tree where
fmap _ Nil = Nil
fmap f (Tip a) = Tip (f a)
fmap f (Bin x y) = Bin (fmap f x) (fmap f y)
empty :: State [a] Bool
empty = liftM null get
item :: State [a] a
item = do (x:xs) <- get
put xs
return x
peekitem :: State [a] a
peekitem = liftM head get
entropy :: State String String
entropy = do ys <- get
case ys of
('\xFF':'\x00':xs) -> do put xs
liftM ('\xFF':) entropy
('\xFF': _ :xs) -> do put xs
entropy
( x :xs) -> do put xs
liftM (x:) entropy
[] -> return []
byte :: State String Int
byte = liftM ord item
word :: State String Int
word = do a<-byte
b<-byte
return $ a*256+b
nibbles :: State String (Int,Int)
nibbles = liftM byte2nibs byte
matrix :: Monad m => Dim -> m a -> m (Mat a)
matrix (y,x) = replicateM y . replicateM x
many :: Monad (State [a]) => State [a] b -> State [a] [b]
many f = do b <- empty
if b
then return []
else liftM2 (:) f (many f)
sf_uncur :: (b -> State a (b,c)) -> State (a,b) c
sf_uncur f = do (a,b) <- get
let g = f b
let ((b2,c),a2) = runState g a
put (a2,b2)
return c
sf_curry :: State (a,b) c -> b -> State a (b,c)
sf_curry (State h) = f
where f b = State g
where g a = ((b2,c),a2)
where (c,(a2,b2)) = h (a,b)
build :: Monad (State [(a,Int)]) => Int -> State [(a,Int)] (Tree a)
build n = do b <- empty
(_,s) <- peekitem
t <- if n==s
then do (v,_) <- item
return $ Tip v
else do x <- build (n+1)
y <- build (n+1)
return $ Bin x y
return $ if b then Nil else t
huffmanTree :: Monad (State [(a,Int)]) => [[a]] -> Tree a
huffmanTree = evalState (build 0) . concat . zipWith f [1..16]
where f s = map (\v->(v,s))
treeLookup :: Tree a -> State Bits a
treeLookup (Tip x) = return x
treeLookup (Bin lef rit) = do b <- item
treeLookup (if b then rit else lef)
treeLookup Nil = error "treeLookup needs nonempty tree"
receive :: Int -> State Bits Int
receive 0 = return 0
receive k = do n <- receive (k1)
b <- item
return $ 2*n + (if b then 1 else 0)
dcdecode :: Tree Int -> State Bits Int
dcdecode t = do s <- treeLookup t
v <- receive s
return $ extend v s
extend :: Int -> Int -> Int
extend v t | t==0 = 0
| v>=vt = v
| otherwise = v + 1 2*vt
where vt = 2^(t1)
acdecode :: Tree (Int,Int) -> Int -> State Bits [Int]
acdecode t k
= do (r,s) <- treeLookup t
let k2 = k + r + 1
if r==0&&s==0
then return (replicate (64k) 0)
else do x <- receive s
xs <- if k2>=64 then return [] else acdecode t k2
return $ replicate r 0 ++ (extend x s:xs)
idct1 :: [Float] -> [Float]
idct1 = matapply cosinuses
idct2 :: Mat Float -> Mat Float
idct2 = transpose . map idct1 . transpose . map idct1
cosinuses :: Mat Float
cosinuses = map f [1,3..15]
where f :: Int -> [Float]
f x = map g [0..7]
where g :: Int -> Float
g 0 = 0.5 / sqrt 2.0
g u = 0.5 * cos(fromIntegral(x*u)*(pi/16.0))
type QuaTab = [Int]
dequant :: QuaTab -> [Int] -> Mat Int8
dequant = matmap truncate `o` idct2 `o` zigzag `o` map fromIntegral `o` zipWith (*)
upsamp :: Dim -> Mat a -> Mat a
upsamp (1,1) = id
upsamp (x,y) = multi y . map (multi x)
zigzag :: [a] -> Mat a
zigzag xs = matmap (xs!!) [[ 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]
]
zigzag2 :: [a] -> Mat a
zigzag2 cs = (transpose . map concat . transpose . fst . foldr f e) [1..15]
where e = ([],reverse cs)
f n (rss,xs) = (bs:rss, ys)
where (as,ys) = splitAt (min n (16n)) xs
rev = if even n then id else reverse
bs = replicate (max (n8) 0) []
++ map (:[]) (rev as)
++ replicate (max (8n) 0) []
type DataUnit = Mat Int8
type Picture = Mat PixelRGB
type DataSpec = (Dim, QuaTab, Tree Int, Tree (Int,Int))
type MCUSpec = [(Dim, DataSpec)]
dataunit :: DataSpec -> Int -> State Bits (Int,DataUnit)
dataunit (u,q,dc,ac) x = do dx <- dcdecode dc
xs <- acdecode ac 1
let y=x+dx
return (y, upsamp u (dequant q (y:xs)))
units :: Dim -> DataSpec -> State (Bits,Int) DataUnit
units dim = fmap matconcat . matrix dim . sf_uncur . dataunit
units2 :: (Dim,DataSpec) -> Int -> State Bits (Int,DataUnit)
units2 = sf_curry . uncurry units
mcu :: MCUSpec -> [ Int -> State Bits (Int,DataUnit) ]
mcu = map units2
mcu2 :: MCUSpec -> [Int] -> [ State Bits (Int,DataUnit) ]
mcu2 = zipWith ($) . mcu
mcu3 :: MCUSpec -> [Int] -> State Bits ([Int],[DataUnit])
mcu3 = fmap unzip `o` sequence `o` mcu2
mcu4 :: MCUSpec -> State (Bits,[Int]) Picture
mcu4 = fmap (matmap yCbCr2RGB . matzip) . sf_uncur . mcu3
picture :: Dim -> MCUSpec -> State (Bits,[Int]) Picture
picture dim = fmap matconcat . matrix dim . mcu4
type FrameCompo = (Int,Dim,Int)
type ScanCompo = (Int,Int,Int)
type QtabCompo = (Int,[Int])
type SOF = (Dim,[FrameCompo])
type DHT = (Int,Int,Tree Int)
type SOS = ([ScanCompo],Bits)
type DQT = [QtabCompo]
type XXX = (Char,String)
frameCompo :: State String (Int, (Int,Int), Int)
frameCompo = do c <- byte
dim <- nibbles
tq <- byte
return $ (c,dim,tq)
scanCompo :: State String (Int,Int,Int)
scanCompo = do cs <- byte
(td,ta) <- nibbles
return $ (cs,td,ta)
qtabCompo :: State String (Int, [Int])
qtabCompo = do (p,ident) <- nibbles
qt <- replicateM 64 (if p==0 then byte else word)
return $ (ident,qt)
sofSeg :: State String ( (Int,Int), [(Int, (Int,Int), Int)] )
sofSeg = do _ <- word
_ <- byte
y <- word
x <- word
n <- byte
fcs <- replicateM n frameCompo
return $ ((y,x), fcs)
dhtSeg :: State String (Int, Int, Tree Int)
dhtSeg = do _ <- word
(tc,th) <- nibbles
ns <- replicateM 16 byte
v <- sequence (map (flip replicateM byte) ns)
return $ (tc, th, huffmanTree v)
dqtSeg :: State String [(Int, [Int])]
dqtSeg = do len <- word
replicateM ((len2)`rem`64) qtabCompo
sosSeg :: State String ( [(Int,Int,Int)], Bits)
sosSeg = do _ <- word
n <- byte
scs <- replicateM n scanCompo
_ <- byte
_ <- byte
_ <- nibbles
ent <- entropy
return $ (scs, string2bits ent)
segment :: (SOF->a, DHT->a, DQT->a, SOS->a, XXX->a) -> State String a
segment (sof,dht,dqt,sos,xxx) =
do _ <- item
c <- item
s <- case c of
'\xC0' -> fmap sof sofSeg
'\xC4' -> fmap dht dhtSeg
'\xDB' -> fmap dqt dqtSeg
'\xDA' -> fmap sos sosSeg
'\xD8' -> return $ xxx (c,[])
'\xD9' -> return $ xxx (c,[])
_ -> do n <- word
s <- replicateM (n2) item
return $ xxx (c,s)
return s
type Huf = (Table(Tree Int), Table(Tree (Int,Int)))
type Sof = (Dim, Table(Dim,QuaTab))
type Qua = Table QuaTab
type State2 = (Sof,Huf,Qua,Picture)
segments :: State String [State2->State2]
segments = many (segment (sof,dht,dqt,sos,xxx))
where sof x s@(_,b,c,d) = (evalSOF x s, b, c, d)
dht x s@(a,_,c,d) = (a, evalDHT x s, c, d)
dqt x s@(a,b,_,d) = (a, b, evalDQT x s, d)
sos x s@(a,b,c,_) = (a, b, c, evalSOS x s)
xxx x s = trace ("extra data: " ++ show x) s
errRes :: State2
errRes = (error "SOF", error "DHT", error "DQT", error "SOS")
evalSOF :: SOF -> State2 -> Sof
evalSOF (dim,xs) (~(_,sof),_,qua,_) = (dim, foldr f sof xs)
where f (i,d,q) = subst i (d,qua q)
evalDHT :: DHT -> State2 -> Huf
evalDHT (0,i,tree) (_,~(hdc,hac),_,_) = (subst i tree hdc, hac)
evalDHT (1,i,tree) (_,~(hdc,hac),_,_) = (hdc, subst i (fmap byte2nibs tree) hac)
evalDHT _ _ = error "evalDHT: unexpected case"
evalDQT :: DQT -> State2 -> Qua
evalDQT xs (_,_,qua,_) = foldr f qua xs
where f (i,q) = subst i q
evalSOS :: SOS -> State2 -> Picture
evalSOS (cs,xs) (((y,x),sof),(h0,h1),_,_)
= map (take x) (take y (evalState thePicture (xs,[0,0,0])))
where thePicture = picture repCount mcuSpec
mcuSpec = map f cs
f (ident,dc,ac)= (d, (upsCount d, qt, h0 dc, h1 ac))
where (d,qt) = sof ident
repCount = ( ceilDiv y (8*maxy), ceilDiv x (8*maxx) )
upsCount (h,w) = ( maxy `div` h, maxx `div` w )
maxy = maximum ( map (fst.fst) mcuSpec )
maxx = maximum ( map (snd.fst) mcuSpec )
jpegDecode :: String -> Picture
jpegDecode = pi4 . foldl (flip ($)) errRes . evalState segments
where pi4 (_,_,_,x) = x
ppmEncode :: Mat PixelRGB -> String
ppmEncode xss
= "P6\n# Creator: Haskell JPEG decoder\n"
++ w ++ " " ++ h ++ "\n255\n"
++ (concat . map rgbPixel2ppmChars . concat) xss
where w = show (length (head xss))
h = show (length xss)
rgbPixel2ppmChars :: PixelRGB -> String
rgbPixel2ppmChars (PixelRGB r g b)
= [ chr (fromIntegral r) , chr (fromIntegral g), chr (fromIntegral b) ]
bmpEncode :: Mat PixelRGB -> String
bmpEncode xss
= bmphead xss
++ concat (map bmpline (reverse xss))
bmphead :: [[a]] -> String
bmphead xss = (concat . map wor )
([ 19778, len, len `div` 65536, 0, 0 ,54, 0, 40
, 0 , w , 0, h, 0 , 1, 24, 0 ] ++ replicate 11 0)
where w = length (head xss)
h = length xss
a = w*3
p = bmpPad a
len = 54 + (a+p)*h
bmpline :: [PixelRGB] -> String
bmpline xs
= let as = concatMap rgbPixel2bmpChars xs
n = bmpPad (length as)
in if n==0
then as
else as ++ replicate n '\0'
bmpPad :: Int -> Int
bmpPad w
= let m = w `mod` 4
in if m==0 then 0 else 4m
rgbPixel2bmpChars :: PixelRGB -> String
rgbPixel2bmpChars (PixelRGB r g b)
= [chr (fromIntegral b), chr (fromIntegral g), chr (fromIntegral r)]
wor :: Int -> String
wor x = [chr (x`rem`256), chr (x`div`256) ]
readBinFile :: String -> IO String
readBinFile f = do h <- openBinaryFile f ReadMode
hGetContents h
writeBinFile :: String -> String -> IO ()
writeBinFile f s = do h <- openBinaryFile f WriteMode
hPutStr h s
hClose h
jpgFile2bmpFile :: String -> String -> IO ()
jpgFile2bmpFile src dst
= do input <- readBinFile src
let output = (bmpEncode . jpegDecode) input
writeBinFile dst output
jpgFile2ppmFile :: String -> String -> IO ()
jpgFile2ppmFile src dst
= do input <- readBinFile src
let output = (ppmEncode . jpegDecode) input
writeBinFile dst output