aboutsummaryrefslogtreecommitdiffstats
path: root/src/PFM.hs
blob: bf805a5ce80dfd6bc894cda81ffebaada75f0438 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
{-|
Module      : PFM
Description : Debevec PFM reader
Copyright   : (c) 2019, Yann Herklotz Grave
License     : GPL-3
Maintainer  : ymherklotz [at] gmail [dot] com
Stability   : experimental
Portability : POSIX

Debevec PFM reader
-}

module PFM where

import           Control.Applicative        ((<|>))
import           Control.Monad              (void)
import           Data.Attoparsec.ByteString (Parser)
import qualified Data.Attoparsec.ByteString as P
import           Data.Binary.Get            (runGet)
import           Data.Binary.IEEE754        (getFloat32be, getFloat32le,
                                             putFloat32be, putFloat32le)
import           Data.Binary.Put            (runPut)
import           Data.ByteString            (ByteString)
import qualified Data.ByteString            as B
import           Data.ByteString.Lazy       (fromStrict, toStrict)
import qualified Data.ByteString.Lazy       as BL
import           Data.Foldable              (fold)
import           Data.Text                  (Text)
import qualified Data.Text                  as T
import qualified Data.Text.Encoding         as T
import           Data.Word                  (Word8)

data Image = Image { width  :: Int
                   , height :: Int
                   , colour :: [Colour]
                   } deriving (Show)

data Colour = Colour { r :: Float
                     , g :: Float
                     , b :: Float
                     }
            | Mono Float
            deriving (Show)

data Endianness = Big | Little

data ImageType = MonoImage | ColourImage

matchText :: Text -> Parser ByteString
matchText = P.string . T.encodeUtf8

magicNum :: Parser ImageType
magicNum = do
  match <- T.decodeUtf8 <$> (matchText "Pf" <|> matchText "PF")
  if match == "Pf"
    then return MonoImage
    else return ColourImage


skipNewline :: Parser ()
skipNewline = P.skip isNewline
  where
    isNewline w = w == 13 || w == 10

skipSpace :: Parser ()
skipSpace = P.skip (== 32)

decode :: (Read a) => [Word8] -> a
decode = read . T.unpack . T.decodeUtf8 . B.pack

matchMult :: String -> Parser [Word8]
matchMult = P.many1 . P.satisfy . P.inClass

num :: Parser Int
num = decode <$> matchMult "0-9"

endianness :: Parser Endianness
endianness =
  getEnd . (<0.0) . decode <$> matchMult "0-9.-"
  where
    getEnd True  = Little
    getEnd False = Big

float :: Endianness -> Parser Float
float e =
  runGet conv . fromStrict <$> P.take 4
  where
    conv = case e of
      Big    -> getFloat32be
      Little -> getFloat32le

header :: Parser (Int, Int, Endianness, ImageType)
header = do
  n <- magicNum
  skipNewline
  n1 <- num
  skipSpace
  n2 <- num
  skipNewline
  s <- endianness
  skipNewline
  return (n1, n2, s, n)

parseColour :: Endianness -> Parser Colour
parseColour e = do
  ri <- float e
  gi <- float e
  bi <- float e
  return $ Colour ri gi bi

parseMono :: Endianness -> Parser Colour
parseMono e = Mono <$> float e

parser :: Parser Image
parser = do
  (w, h, e, i) <- header
  c <- P.many1 $ fun i e
  return $ Image w h c
  where
    fun i = case i of
      ColourImage -> parseColour
      MonoImage   -> parseMono

imageType :: [Colour] -> Text
imageType []           = "PF"
imageType (Colour{}:_) = "PF"
imageType (Mono{}:_)   = "Pf"

tShow :: (Show a) => a -> Text
tShow = T.pack . show

encFloat :: Float -> BL.ByteString
encFloat = runPut . putFloat32le

encodeColour :: Colour -> BL.ByteString
encodeColour (Colour ri gi bi) =
  encFloat ri <> encFloat gi <> encFloat bi
encodeColour (Mono m) =
  encFloat m

encode :: Image -> BL.ByteString
encode (Image w h c) =
  fromStrict (T.encodeUtf8 he) <> body
  where
    he = imageType c <> "\n" <> tShow w <> " " <> tShow h <> "\n-1.0\n"
    body = fold $ encodeColour <$> c

parse :: ByteString -> Image
parse s = case P.parseOnly parser s of
  Left str -> error str
  Right i  -> i