src/Volumetric.hs

   1 {-# LANGUAGE FlexibleContexts #-}
   2
   3 -- | The Volumetric module contains functions for manipulating the
   4 --   volumetric data files found at,
   5 --
   6 --     <http://graphics.stanford.edu/data/voldata/>
   7 --
   8 module Volumetric (
   9   bracket_array,
  10   flip_x,
  11   flip_y,
  12   read_word16s,
  13   round_array,
  14   swap_bytes,
  15   write_values_to_bmp,
  16   write_word16s,
  17   z_slice )
  18 where
  19
  20 import Data.Word ( Word8, Word16 )
  21 import Data.Bits ( (.&.), (.|.), shift )
  22 import qualified Data.Array.Repa as Repa ( map )
  23 import Data.Array.Repa (
  24   (:.)( (:.) ),
  25   Array,
  26   D,
  27   DIM2,
  28   DIM3,
  29   Shape,
  30   Source,
  31   U,
  32   Z( Z ),
  33   computeUnboxedP,
  34   copyP,
  35   slice )
  36 import Data.Array.Repa.Algorithms.ColorRamp ( rampColorHotToCold )
  37 import Data.Array.Repa.Eval ( now )
  38 import Data.Array.Repa.IO.Binary (
  39   readArrayFromStorableFile,
  40   writeArrayToStorableFile )
  41 import Data.Array.Repa.IO.BMP ( writeImageToBMP )
  42 import Data.Array.Repa.Operators.Traversal ( unsafeTraverse )
  43 import Data.Array.Repa.Repr.Unboxed ( Unbox )
  44 import Data.Array.Repa.Slice ( All( All ), Any( Any ) )
  45
  46 import Values ( Values, Values2D )
  47
  48
  49 -- | RawData is an array of words (16 bits), as contained in the
  50 --   volumetric data files.
  51 type RawData sh = Array U sh Word16
  52
  53 -- | A specialization of the 'RawData' type, to three dimensions.
  54 type RawData3D = RawData DIM3
  55
  56 type RGB = (Word8, Word8, Word8)
  57 type ColorData sh = Array U sh RGB
  58
  59
  60 {-# INLINE read_word16s #-}
  61 read_word16s :: FilePath -> DIM3 -> IO RawData3D
  62 read_word16s path shape = do
  63   arr <- readArrayFromStorableFile path shape
  64   c   <- copyP arr
  65   now c
  66
  67
  68
  69 bracket :: Double -> Double -> Double -> Word16
  70 bracket lower_threshold upper_threshold x
  71         | x < lower_threshold      = 0
  72         | x > upper_threshold      = 255
  73         | otherwise                = truncate (r * 255)
  74             where
  75               numerator = x - lower_threshold
  76               denominator = upper_threshold - lower_threshold
  77               r = numerator/denominator
  78
  79
  80 flip16 :: Word16 -> Word16
  81 flip16 xx =
  82   shift xx 8 .|. (shift xx (-8) .&. 0x00ff)
  83
  84
  85 {-# INLINE swap_bytes #-}
  86 swap_bytes :: (Shape sh, Source r Word16) => Array r sh Word16
  87              -> Array D sh Word16
  88 swap_bytes =
  89   Repa.map flip16
  90
  91
  92 bracket_array :: Shape sh => Double -> Double -> Values sh -> Array D sh Word16
  93 bracket_array lt ut =
  94   Repa.map (bracket lt ut)
  95
  96
  97 {-# INLINE round_array #-}
  98 round_array :: Shape sh => Values sh -> Array D sh Word16
  99 round_array =
 100   Repa.map round
 101
 102
 103 flip_y :: Source r Word16 => Int -> Array r DIM3 Word16 -> Array D DIM3 Word16
 104 flip_y height arr =
 105   unsafeTraverse arr id
 106               (\get (Z :. z :. y :. x) ->
 107                    get (Z :. z :. (height - 1) - y :. x))
 108
 109 flip_x :: Source r Word16 => Int -> Array r DIM3 Word16 -> Array D DIM3 Word16
 110 flip_x width arr =
 111   unsafeTraverse arr id
 112               (\get (Z :. z :. y :. x) ->
 113                    get (Z :. z :. y :. (width - 1) - x))
 114
 115
 116 {-# INLINE write_word16s #-}
 117 write_word16s :: (Shape sh) => FilePath -> (RawData sh) -> IO ()
 118 write_word16s = writeArrayToStorableFile
 119
 120
 121
 122 --
 123 -- Instead of IO, we could get away with a generic monad 'm'
 124 -- here. However, /we/ only call this function from within IO.
 125 --
 126 values_to_colors :: (Shape sh) => (Values sh) -> IO (ColorData sh)
 127 values_to_colors arr =
 128   computeUnboxedP $ Repa.map (truncate_rgb . ramp_it) arr
 129   where
 130     ramp_it :: Double -> (Double, Double, Double)
 131     ramp_it x =
 132       if x == 0 then
 133         (0, 0, 0)
 134       else
 135         rampColorHotToCold 0 255 x
 136
 137     truncate_rgb :: (Double, Double, Double) -> (Word8, Word8, Word8)
 138     truncate_rgb (r, g, b) =
 139       (r', g', b')
 140       where
 141         r' = truncate (r * 255)
 142         g' = truncate (g * 255)
 143         b' = truncate (b * 255)
 144
 145
 146 write_values_to_bmp :: FilePath -> Values2D -> IO ()
 147 write_values_to_bmp path values = do
 148   colors <- values_to_colors values
 149   writeImageToBMP path colors
 150
 151
 152 z_slice :: (Unbox a, Source r a) => Int -> Array r DIM3 a -> Array D DIM2 a
 153 z_slice n arr =
 154   slice arr (Any :. n :. All :. All)