Finished coursework

3 files changed, 81 insertions, 9 deletions

diff --git a/mirror-ball.cabal b/mirror-ball.cabal
index 44a0b9c..314c440 100644
--- a/mirror-ball.cabal
+++ b/mirror-ball.cabal
@@ -17,8 +17,11 @@ executable mirror-ball
   hs-source-dirs:      src
   main-is:             Main.hs
   default-language:    Haskell2010
+  other-modules:       Vec
+  ghc-options:         -threaded
   build-depends:       base >= 4.7 && < 5
                      , bytestring >= 0.10.8.2
+                     , criterion >= 1.5.3.0
                      , directory >= 1.3.3.0
                      , pfm >= 0.1.0.0
                      , text >= 1.2.3.1
diff --git a/src/Main.hs b/src/Main.hs
index 8bba9a0..a0b3889 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -1,5 +1,7 @@
 module Main where
 
+import           Control.Concurrent
+import           Criterion
 import           Data.Bifunctor       (Bifunctor, bimap)
 import qualified Data.ByteString      as B
 import qualified Data.ByteString.Lazy as BL
@@ -30,17 +32,20 @@ uncurry3 f (a, b, c) = f a b c
 colourCorrect :: Double -> Double
 colourCorrect = (+0.5) . (/2)
 
-toCircle :: ((Int, Int) -> Vec Double) -> Int -> [[(Int, Int)]] -> [[PFMColour]]
+toCircle :: ((Int, Int) -> PFMColour) -> Int -> [[(Int, Int)]] -> [[PFMColour]]
 toCircle f v m =
   fmap toCol <$> m
   where
     toCol val = if dist val <= v ^ (2 :: Int)
-      then uncurry3 PFMColour (unVec $ realToFrac . colourCorrect <$> f val)
+      then f val
       else PFMColour 0 0 0--0.5 0.5 1.0
     dist = add . bmap ((^(2 :: Int)) . (v-))
     add (an, bn) = an + bn
 
-circleImage :: ((Int, Int) -> Vec Double) -> Int -> PFMImage
+toColour :: ((Int, Int) -> Vec Double) -> (Int, Int) -> PFMColour
+toColour f val = uncurry3 PFMColour (unVec $ realToFrac . colourCorrect <$> f val)
+
+circleImage :: ((Int, Int) -> PFMColour) -> Int -> PFMImage
 circleImage f s = PFMImage s s $ toCircle f ((s - 1) `div` 2)
   [ [ (y, x) | x <- [0..s-1] ] | y <- reverse [0..s-1] ]
 
@@ -48,12 +53,53 @@ revColour :: PFMImage -> PFMImage
 revColour (PFMImage w h i) =
   PFMImage w h $ reverse i
 
+bench :: IO ()
+bench = do
+  im <- B.readFile "data/urbanEM_latlong.pfm"
+  let urban = revColour $ parse im
+  benchmark $ whnf (circleImage (toColour $ normalise 511)) 511
+  benchmark $ whnf (circleImage (toColour $ reflect 511 (Vec (0, 0, 1)))) 511
+  benchmark $ whnf (circleImage (fromImage 511 (Vec (0, 0, 1)) urban)) 511
+
+indexTwice :: [[a]] -> (Int, Int) -> a
+indexTwice c (y, x) = c !! y !! x
+
+fromImage :: Int -> Vec Double -> PFMImage -> (Int, Int) -> PFMColour
+fromImage s v (PFMImage w h c) (y, x) =
+  indexTwice c . indexLatLong w h . toSpherical $ reflect s v (y, x)
+
+applyGamma :: Float -> PFMImage -> PFMImage
+applyGamma g (PFMImage w h c) =
+  PFMImage w h $ fmap gc <$> c
+  where
+    gc (PFMColour r gr b) = PFMColour (gamma g r) (gamma g gr) (gamma g b)
+
+gamma :: (Floating a) => a -> a -> a
+gamma g m = m ** (1 / g)
+
+myForkIO :: IO () -> IO (MVar ())
+myForkIO io = do
+  mvar <- newEmptyMVar
+  _ <- forkFinally io (\_ -> putMVar mvar ())
+  return mvar
+
 main :: IO ()
 main = do
+  im <- B.readFile "data/urbanEM_latlong.pfm"
+  let urban = revColour $ parse im
+  let i = circleImage (toColour $ normalise 511) 511
+  let r = circleImage (toColour $ reflect 511 (Vec (0, 0, 1))) 511
+  let f = circleImage (fromImage 511 (Vec (0, 0, 1)) urban) 511
+  let fg = applyGamma 1.7 f
   createDirectoryIfMissing True "data"
-  let i = circleImage (normalise 511) 511
-  BL.writeFile "data/normal.ppm" . encodePPM . clampImage $ i
-  BL.writeFile "data/normal.pfm" . encode . revColour $ i
-  let r = circleImage (reflect 511 (Vec (0, 0, 1))) 511
-  BL.writeFile "data/reflect.ppm" . encodePPM . clampImage $ r
-  BL.writeFile "data/reflect.pfm" . encode . revColour $ r
+  vars <- sequence $ myForkIO <$>
+    [ BL.writeFile "data/normal.ppm" . encodePPM . clampImage $ i
+    , BL.writeFile "data/normal.pfm" . encode . revColour $ i
+    , BL.writeFile "data/reflect.ppm" . encodePPM . clampImage $ r
+    , BL.writeFile "data/reflect.pfm" . encode . revColour $ r
+    , BL.writeFile "data/final_nogamma.ppm" . encodePPM . clampImage $ f
+    , BL.writeFile "data/final_nogamma.pfm" . encode . revColour $ f
+    , BL.writeFile "data/final.ppm" . encodePPM . clampImage $ fg
+    , BL.writeFile "data/final.pfm" . encode . revColour $ fg
+    ]
+  sequence_ $ takeMVar <$> vars
diff --git a/src/Vec.hs b/src/Vec.hs
index 9e4fdbd..3f5a839 100644
--- a/src/Vec.hs
+++ b/src/Vec.hs
@@ -26,6 +26,12 @@ instance (Num a) => Num (Vec a) where
   signum = fmap signum
   fromInteger i = Vec (fromInteger i, 0, 0)
 
+newtype Sph a = Sph { unSph :: (a, a) }
+              deriving (Show)
+
+instance Functor Sph where
+  fmap f (Sph (a, b)) = Sph (f a, f b)
+
 findZ :: (Floating a) => a -> a -> Vec a
 findZ x y =
   Vec (x, y, sqrt (1 - x^^2 - y^^2))
@@ -46,3 +52,20 @@ reflect size v (y, x) =
   where
     n = normalise size (y, x)
     l = ((2 * dot n v)*) <$> n
+
+toSpherical :: (Floating a, Eq a, Ord a) => Vec a -> Sph a
+toSpherical (Vec (x, y, z))
+  | z == 0 && x >= 0 =
+    Sph (acos y, pi / 2)
+  | z == 0 =
+    Sph (acos y, - pi / 2)
+  | z < 0 =
+    Sph (acos y, signum x * pi + atan (x / z))
+  | otherwise =
+    Sph (acos y, atan (x / z))
+
+indexLatLong :: (RealFrac a, Floating a) => Int -> Int -> Sph a -> (Int, Int)
+indexLatLong w h (Sph (theta, phi)) =
+  ( floor $ theta / pi * fromIntegral h
+  , floor $ ((phi / (2 * pi)) + 0.5) * fromIntegral w
+  )