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 MRI (
  18   flip_x,
  19   flip_y,
  20   read_word16s,
  21   round_array,
  22   swap_bytes,
  23   write_values_slice_to_bitmap,
  24   write_word16s,
  25   z_slice
  26   )
  27
  28
  29 validate_args :: Args -> IO ()
  30 validate_args Args{..} = do
  31   when (scale <= 0) $ do
  32     hPutStrLn stderr "ERROR: scale must be greater than zero."
  33     exitWith (ExitFailure exit_arg_not_positive)
  34
  35   when (width <= 0) $ do
  36     hPutStrLn stderr "ERROR: width must be greater than zero."
  37     exitWith (ExitFailure exit_arg_not_positive)
  38
  39   when (height <= 0) $ do
  40     hPutStrLn stderr "ERROR: height must be greater than zero."
  41     exitWith (ExitFailure exit_arg_not_positive)
  42
  43   when (depth <= 0) $ do
  44     hPutStrLn stderr "ERROR: depth must be greater than zero."
  45     exitWith (ExitFailure exit_arg_not_positive)
  46
  47   case slice of
  48     Just s ->
  49       when (s < 0 || s > depth) $ do
  50         hPutStrLn stderr "ERROR: slice must be between zero and depth."
  51         exitWith (ExitFailure exit_arg_out_of_bounds)
  52     Nothing -> return ()
  53
  54
  55 main :: IO ()
  56 main = do
  57   args@Args{..} <- apply_args
  58   -- validate_args will simply exit if there's a problem.
  59   validate_args args
  60
  61   -- The first thing we do is set the number of processors. We get the
  62   -- number of processors (cores) in the machine with
  63   -- getNumProcessors, and set it with setNumCapabilities. This is so
  64   -- we don't have to pass +RTS -Nfoo on the command line every time.
  65   num_procs <- getNumProcessors
  66   setNumCapabilities num_procs
  67
  68   -- Determine whether we're doing 2d or 3d. If we're given a slice,
  69   -- assume 2d.
  70   let mri_shape = (R.Z R.:. depth R.:. height R.:. width) :: R.DIM3
  71
  72   if (isJust slice) then
  73     main2d args mri_shape
  74   else
  75     main3d args mri_shape
  76
  77   exitSuccess
  78
  79   where
  80
  81
  82
  83 main3d :: Args -> R.DIM3 -> IO ()
  84 main3d Args{..} mri_shape = do
  85   let zoom_factor = (scale, scale, scale)
  86   arr <- read_word16s input mri_shape
  87   let arr'          = swap_bytes arr
  88   let arrMRI        = R.reshape mri_shape arr'
  89   dbl_data <- R.computeUnboxedP $ R.map fromIntegral arrMRI
  90   raw_output <- zoom dbl_data zoom_factor
  91   word16_output <- R.computeUnboxedP $ round_array raw_output
  92   write_word16s output word16_output
  93
  94
  95 main2d :: Args -> R.DIM3 -> IO ()
  96 main2d Args{..} mri_shape = do
  97   let zoom_factor = (1, scale, scale)
  98   arr <- read_word16s input mri_shape
  99   arrSlice <- R.computeUnboxedP
 100                $ z_slice (fromJust slice)
 101                $ flip_x width
 102                $ flip_y height
 103                $ swap_bytes arr
 104   let arrSlice' = R.reshape mri_slice3d arrSlice
 105
 106   -- If zoom isn't being inlined we need to extract the slice before hand,
 107   -- and convert it to the require formed.
 108   dbl_data      <- R.computeUnboxedP $ R.map fromIntegral arrSlice'
 109   raw_output    <- zoom dbl_data zoom_factor
 110   arrSlice0     <- R.computeUnboxedP $ z_slice 0 raw_output
 111
 112   write_values_slice_to_bitmap arrSlice0 output
 113   where
 114     mri_slice3d :: R.DIM3
 115     mri_slice3d = (R.Z R.:. 1 R.:. height  R.:. width)