src/Main.hs

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