src/Main.hs

   1 {-# LANGUAGE RecordWildCards, DoAndIfThenElse #-}
   2
   3 module Main
   4 where
   5
   6 import Data.Maybe (fromJust)
   7 import Control.Monad (when)
   8 import qualified Data.Array.Repa as R
   9 import Data.Maybe (isJust)
  10 import GHC.Conc (getNumProcessors, setNumCapabilities)
  11 import System.IO (hPutStrLn, stderr)
  12 import System.Exit (exitSuccess, exitWith, ExitCode(..))
  13
  14 import CommandLine (Args(..), apply_args)
  15 import ExitCodes
  16 import Grid (zoom)
  17 import Volumetric (
  18   bracket_array,
  19   flip_x,
  20   flip_y,
  21   read_word16s,
  22   round_array,
  23   swap_bytes,
  24   write_values_to_bmp,
  25   write_word16s,
  26   z_slice
  27   )
  28
  29
  30 validate_args :: Args -> IO ()
  31 validate_args Args{..} = do
  32   when (scale <= 0) $ do
  33     hPutStrLn stderr "ERROR: scale must be greater than zero."
  34     exitWith (ExitFailure exit_arg_not_positive)
  35
  36   when (width <= 0) $ do
  37     hPutStrLn stderr "ERROR: width must be greater than zero."
  38     exitWith (ExitFailure exit_arg_not_positive)
  39
  40   when (height <= 0) $ do
  41     hPutStrLn stderr "ERROR: height must be greater than zero."
  42     exitWith (ExitFailure exit_arg_not_positive)
  43
  44   when (depth <= 0) $ do
  45     hPutStrLn stderr "ERROR: depth must be greater than zero."
  46     exitWith (ExitFailure exit_arg_not_positive)
  47
  48   case slice of
  49     Just s ->
  50       when (s < 0 || s > depth) $ do
  51         hPutStrLn stderr "ERROR: slice must be between zero and depth."
  52         exitWith (ExitFailure exit_arg_out_of_bounds)
  53     Nothing -> return ()
  54
  55
  56 main :: IO ()
  57 main = do
  58   args@Args{..} <- apply_args
  59   -- validate_args will simply exit if there's a problem.
  60   validate_args args
  61
  62   -- The first thing we do is set the number of processors. We get the
  63   -- number of processors (cores) in the machine with
  64   -- getNumProcessors, and set it with setNumCapabilities. This is so
  65   -- we don't have to pass +RTS -Nfoo on the command line every time.
  66   num_procs <- getNumProcessors
  67   setNumCapabilities num_procs
  68
  69   -- Determine whether we're doing 2d or 3d. If we're given a slice,
  70   -- assume 2d.
  71   let shape = (R.Z R.:. depth R.:. height R.:. width) :: R.DIM3
  72
  73   if (isJust slice) then
  74     main2d args shape
  75   else
  76     main3d args shape
  77
  78   exitSuccess
  79
  80   where
  81
  82
  83
  84 main3d :: Args -> R.DIM3 -> IO ()
  85 main3d Args{..} shape = do
  86   let zoom_factor = (scale, scale, scale)
  87   arr <- read_word16s input shape
  88   let arr_swapped = swap_bytes arr
  89   let arr_shaped  = R.reshape shape arr_swapped
  90   dbl_data <- R.computeUnboxedP $ R.map fromIntegral arr_shaped
  91   raw_output <- zoom dbl_data zoom_factor
  92   let word16_output = round_array raw_output
  93   -- Switch the bytes order back to what it was. This lets us use the
  94   -- same program to view the input/output data.
  95   swapped_output <- R.computeUnboxedP $ swap_bytes word16_output
  96   write_word16s output swapped_output
  97
  98
  99 main2d :: Args -> R.DIM3 -> IO ()
 100 main2d Args{..} shape = do
 101   let zoom_factor = (1, scale, scale)
 102   arr <- read_word16s input shape
 103   arrSlice <- R.computeUnboxedP
 104                $ z_slice (fromJust slice)
 105                $ flip_x width
 106                $ flip_y height
 107                $ swap_bytes arr
 108   let arrSlice' = R.reshape slice3d arrSlice
 109
 110   -- If zoom isn't being inlined we need to extract the slice before hand,
 111   -- and convert it to the require formed.
 112   dbl_data   <- R.computeUnboxedP $ R.map fromIntegral arrSlice'
 113   raw_output <- zoom dbl_data zoom_factor
 114   arrSlice0  <- R.computeUnboxedP $ z_slice 0 raw_output
 115
 116   -- Make doubles from the thresholds which are given as Ints.
 117   let lt = fromIntegral lower_threshold
 118   let ut = fromIntegral upper_threshold
 119
 120   let arr_bracketed = bracket_array lt ut arrSlice0
 121   values <- R.computeUnboxedP $ R.map fromIntegral arr_bracketed
 122   write_values_to_bmp output values
 123
 124   where
 125     slice3d :: R.DIM3
 126     slice3d = (R.Z R.:. 1 R.:. height  R.:. width)