Move the last two MRI-specific variables out of MRI.hs and into the command-line...

[spline3.git] / src / Volumetric.hs

diff --git a/src/Volumetric.hs b/src/Volumetric.hs
new file mode 100644 (file)
index 0000000..24bc13b
--- /dev/null
+++ b/src/Volumetric.hs
@@ -0,0 +1,137 @@
+{-# LANGUAGE FlexibleContexts #-}
+-- | The Volumetric module contains functions for manipulating the
+--   volumetric data files found at,
+--
+--     <http://graphics.stanford.edu/data/voldata/>
+--
+module Volumetric (
+  bracket_array,
+  flip_x,
+  flip_y,
+  read_word16s,
+  round_array,
+  swap_bytes,
+  write_values_to_bmp,
+  write_word16s,
+  z_slice
+  )
+where
+
+import Data.Word
+import Data.Bits
+import Data.Array.Repa                          as R
+import Data.Array.Repa.Eval                     as R (now)
+import Data.Array.Repa.Repr.Unboxed             as R
+import Data.Array.Repa.IO.Binary                as R
+import Data.Array.Repa.Algorithms.ColorRamp     as R
+import Data.Array.Repa.Operators.Traversal      as R (unsafeTraverse)
+import Data.Array.Repa.IO.BMP                   as R (writeImageToBMP)
+
+import Values
+
+-- | RawData is an array of words (16 bits), as contained in the
+--   volumetric data files.
+type RawData sh = Array U sh Word16
+
+-- | A specialization of the 'RawData' type, to three dimensions.
+type RawData3D = RawData DIM3
+
+type RGB = (Word8, Word8, Word8)
+type ColorData sh = Array U sh RGB
+
+
+{-# INLINE read_word16s #-}
+read_word16s :: FilePath -> DIM3 -> IO RawData3D
+read_word16s path shape = do
+  arr <- R.readArrayFromStorableFile path shape
+  c   <- R.copyP arr
+  now $ c
+
+
+
+bracket :: Double -> Double -> Double -> Word16
+bracket lower_threshold upper_threshold x
+        | x < lower_threshold      = 0
+        | x > upper_threshold      = 255
+        | otherwise                    = truncate (r * 255)
+            where
+              numerator = x - lower_threshold
+              denominator = upper_threshold - lower_threshold
+              r = numerator/denominator
+
+
+flip16 :: Word16 -> Word16
+flip16 xx =
+  shift xx 8 .|. (shift xx (-8) .&. 0x00ff)
+
+
+{-# INLINE swap_bytes #-}
+swap_bytes :: (Shape sh, Source r Word16) => Array r sh Word16
+             -> Array D sh Word16
+swap_bytes =
+  R.map flip16
+
+
+bracket_array :: Shape sh => Double -> Double -> Values sh -> Array D sh Word16
+bracket_array lt ut =
+  R.map (bracket lt ut)
+
+
+{-# INLINE round_array #-}
+round_array :: Shape sh => Values sh -> Array D sh Word16
+round_array =
+  R.map round
+
+
+flip_y :: Source r Word16 => Int -> Array r DIM3 Word16 -> Array D DIM3 Word16
+flip_y height arr =
+  R.unsafeTraverse arr id
+              (\get (Z :. z :. y :. x) ->
+                   get (Z :. z :. (height - 1) - y :. x))
+
+flip_x :: Source r Word16 => Int -> Array r DIM3 Word16 -> Array D DIM3 Word16
+flip_x width arr =
+  R.unsafeTraverse arr id
+              (\get (Z :. z :. y :. x) ->
+                   get (Z :. z :. y :. (width - 1) - x))
+
+
+{-# INLINE write_word16s #-}
+write_word16s :: (Shape sh) => FilePath -> (RawData sh) -> IO ()
+write_word16s = R.writeArrayToStorableFile
+
+
+
+--
+-- Instead of IO, we could get away with a generic monad 'm'
+-- here. However, /we/ only call this function from within IO.
+--
+values_to_colors :: (Shape sh) => (Values sh) -> IO (ColorData sh)
+values_to_colors arr =
+  R.computeUnboxedP $ R.map (truncate_rgb . ramp_it) arr
+  where
+    ramp_it :: Double -> (Double, Double, Double)
+    ramp_it x =
+      if x == 0 then
+        (0, 0, 0)
+      else
+        rampColorHotToCold 0 255 x
+
+    truncate_rgb :: (Double, Double, Double) -> (Word8, Word8, Word8)
+    truncate_rgb (r, g, b) =
+      (r', g', b')
+      where
+        r' = truncate (r * 255)
+        g' = truncate (g * 255)
+        b' = truncate (b * 255)
+
+
+write_values_to_bmp :: FilePath -> Values2D -> IO ()
+write_values_to_bmp path values = do
+  colors <- values_to_colors values
+  R.writeImageToBMP path colors
+
+
+z_slice :: (R.Unbox a, Source r a) => Int -> Array r DIM3 a -> Array D DIM2 a
+z_slice n arr =
+  slice arr (Any :. n :. All :. All)