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.Slice ( All( All ), Any( Any ) )
  44
  45 import Values ( Values, Values2D )
  46
  47
  48 -- | RawData is an array of words (16 bits), as contained in the
  49 --   volumetric data files.
  50 type RawData sh = Array U sh Word16
  51
  52 -- | A specialization of the 'RawData' type, to three dimensions.
  53 type RawData3D = RawData DIM3
  54
  55 type RGB = (Word8, Word8, Word8)
  56 type ColorData sh = Array U sh RGB
  57
  58
  59 {-# INLINE read_word16s #-}
  60 read_word16s :: FilePath -> DIM3 -> IO RawData3D
  61 read_word16s path shape = do
  62   arr <- readArrayFromStorableFile path shape
  63   c   <- copyP arr
  64   now c
  65
  66
  67
  68 bracket :: Double -> Double -> Double -> Word16
  69 bracket lower_threshold upper_threshold x
  70         | x < lower_threshold      = 0
  71         | x > upper_threshold      = 255
  72         | otherwise                = truncate (r * 255)
  73             where
  74               numerator = x - lower_threshold
  75               denominator = upper_threshold - lower_threshold
  76               r = numerator/denominator
  77
  78
  79 flip16 :: Word16 -> Word16
  80 flip16 xx =
  81   shift xx 8 .|. (shift xx (-8) .&. 0x00ff)
  82
  83
  84 {-# INLINE swap_bytes #-}
  85 swap_bytes :: (Shape sh, Source r Word16) => Array r sh Word16
  86              -> Array D sh Word16
  87 swap_bytes =
  88   Repa.map flip16
  89
  90
  91 bracket_array :: Shape sh => Double -> Double -> Values sh -> Array D sh Word16
  92 bracket_array lt ut =
  93   Repa.map (bracket lt ut)
  94
  95
  96 {-# INLINE round_array #-}
  97 round_array :: Shape sh => Values sh -> Array D sh Word16
  98 round_array =
  99   Repa.map round
 100
 101
 102 flip_y :: Source r Word16 => Int -> Array r DIM3 Word16 -> Array D DIM3 Word16
 103 flip_y height arr =
 104   unsafeTraverse arr id
 105               (\get (Z :. z :. y :. x) ->
 106                    get (Z :. z :. (height - 1) - y :. x))
 107
 108 flip_x :: Source r Word16 => Int -> Array r DIM3 Word16 -> Array D DIM3 Word16
 109 flip_x width arr =
 110   unsafeTraverse arr id
 111               (\get (Z :. z :. y :. x) ->
 112                    get (Z :. z :. y :. (width - 1) - x))
 113
 114
 115 {-# INLINE write_word16s #-}
 116 write_word16s :: (Shape sh) => FilePath -> (RawData sh) -> IO ()
 117 write_word16s = writeArrayToStorableFile
 118
 119
 120
 121 --
 122 -- Instead of IO, we could get away with a generic monad 'm'
 123 -- here. However, /we/ only call this function from within IO.
 124 --
 125 values_to_colors :: (Shape sh) => (Values sh) -> IO (ColorData sh)
 126 values_to_colors arr =
 127   computeUnboxedP $ Repa.map (truncate_rgb . ramp_it) arr
 128   where
 129     ramp_it :: Double -> (Double, Double, Double)
 130     ramp_it x =
 131       if x == 0 then
 132         (0, 0, 0)
 133       else
 134         rampColorHotToCold 0 255 x
 135
 136     truncate_rgb :: (Double, Double, Double) -> (Word8, Word8, Word8)
 137     truncate_rgb (r, g, b) =
 138       (r', g', b')
 139       where
 140         r' = truncate (r * 255) :: Word8
 141         g' = truncate (g * 255) :: Word8
 142         b' = truncate (b * 255) :: Word8
 143
 144
 145 write_values_to_bmp :: FilePath -> Values2D -> IO ()
 146 write_values_to_bmp path values = do
 147   colors <- values_to_colors values
 148   writeImageToBMP path colors
 149
 150
 151 z_slice :: (Source r a) => Int -> Array r DIM3 a -> Array D DIM2 a
 152 z_slice n arr =
 153   slice arr (Any :. n :. All :. All)