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