X-Git-Url: http://gitweb.michael.orlitzky.com/?a=blobdiff_plain;f=src%2FMain.hs;h=ab37ca558f009bd30d5ac610e117f9282836442f;hb=3d25d1792dffeee4e73030cdb54ebc12b89afbd2;hp=e577351cc257615a47b60ac455cb2376f2527e5c;hpb=3f0b6b7faecc561af0b7312a11c73a44a1b416f6;p=spline3.git

diff --git a/src/Main.hs b/src/Main.hs
index e577351..ab37ca5 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -1,47 +1,124 @@
-module Main
+{-# LANGUAGE RecordWildCards, DoAndIfThenElse #-}
+
+module Main (main)
 where
 
+import Control.Monad ( when )
 import qualified Data.Array.Repa as R
-import System.Environment (getArgs)
+import Data.Maybe ( fromJust )
+import GHC.Conc ( getNumProcessors, setNumCapabilities )
+import System.IO ( hPutStrLn, stderr )
+import System.Exit (
+  ExitCode( ExitFailure ),
+  exitSuccess,
+  exitWith )
+
+import CommandLine ( Args(..), apply_args )
+import ExitCodes ( exit_arg_not_positive, exit_arg_out_of_bounds )
+import Grid ( zoom )
+import Volumetric (
+  bracket_array,
+  flip_x,
+  flip_y,
+  read_word16s,
+  round_array,
+  swap_bytes,
+  write_values_to_bmp,
+  write_word16s,
+  z_slice )
+
+
+validate_args :: Args -> IO ()
+validate_args Args{..} = do
+  when (scale <= 0) $ do
+    hPutStrLn stderr "ERROR: scale must be greater than zero."
+    exitWith (ExitFailure exit_arg_not_positive)
 
-import Grid (zoom)
-import MRI
+  when (width <= 0) $ do
+    hPutStrLn stderr "ERROR: width must be greater than zero."
+    exitWith (ExitFailure exit_arg_not_positive)
 
+  when (height <= 0) $ do
+    hPutStrLn stderr "ERROR: height must be greater than zero."
+    exitWith (ExitFailure exit_arg_not_positive)
 
-in_file :: FilePath
-in_file = "./data/mri.bin"
+  when (depth <= 0) $ do
+    hPutStrLn stderr "ERROR: depth must be greater than zero."
+    exitWith (ExitFailure exit_arg_not_positive)
+
+  case slice of
+    Just s ->
+      when (s < 0 || s > depth) $ do
+        hPutStrLn stderr "ERROR: slice must be between zero and depth."
+        exitWith (ExitFailure exit_arg_out_of_bounds)
+    Nothing -> return ()
 
 
 main :: IO ()
-main = main3d
+main = do
+  args@Args{..} <- apply_args
+  -- validate_args will simply exit if there's a problem.
+  validate_args args
+
+  -- The first thing we do is set the number of processors. We get the
+  -- number of processors (cores) in the machine with
+  -- getNumProcessors, and set it with setNumCapabilities. This is so
+  -- we don't have to pass +RTS -Nfoo on the command line every time.
+  num_procs <- getNumProcessors
+  setNumCapabilities num_procs
+
+  let shape = (R.Z R.:. depth R.:. height R.:. width) :: R.DIM3
+
+  -- Determine whether we're doing 2d or 3d. If we're given a slice,
+  -- assume 2d.
+  let main_function = case slice of
+                        Nothing -> main3d
+                        Just _  -> main2d
 
+  main_function args shape
+  exitSuccess
 
-main3d :: IO ()
-main3d = do
-  (s:_) <- getArgs
-  let scale = read s :: Int
+
+main3d :: Args -> R.DIM3 -> IO ()
+main3d Args{..} shape = do
   let zoom_factor = (scale, scale, scale)
-  let out_file = "output.bin"
-  arr <- read_word16s in_file
-  let arr' = swap_bytes arr
-  let arrMRI = R.reshape mri_shape arr'
-  let dbl_data = R.force $ R.map fromIntegral arrMRI
-  let output = zoom dbl_data zoom_factor
-  let word16_output = bracket_array output
-  write_word16s out_file word16_output
-
-
-main2d :: IO ()
-main2d = do
-  (s:_) <- getArgs
-  let scale = read s :: Int
-  let zoom_factor = (1, scale, scale)
-  let out_file = "output.bmp"
-  arr <- read_word16s in_file
-  let arr' = swap_bytes arr
-  let arrInv   = flip_x $ flip_y arr'
-  let arrSlice = z_slice 50 arrInv
-  let arrSlice' = R.reshape mri_slice3d arrSlice
-  let dbl_data = R.map fromIntegral arrSlice'
-  let output = zoom dbl_data zoom_factor
-  write_values_slice_to_bitmap (z_slice 0 output) out_file
+  arr <- read_word16s input shape
+  let arr_swapped = swap_bytes arr
+  let arr_shaped  = R.reshape shape arr_swapped
+  dbl_data <- R.computeUnboxedP $ R.map fromIntegral arr_shaped
+  raw_output <- zoom dbl_data zoom_factor
+  let word16_output = round_array raw_output
+  -- Switch the bytes order back to what it was. This lets us use the
+  -- same program to view the input/output data.
+  swapped_output <- R.computeUnboxedP $ swap_bytes word16_output
+  write_word16s output swapped_output
+
+
+main2d :: Args -> R.DIM3 -> IO ()
+main2d Args{..} shape = do
+  let zoom_factor = (1 :: Int, scale, scale)
+  arr <- read_word16s input shape
+  arrSlice <- R.computeUnboxedP
+               $ z_slice (fromJust slice)
+               $ flip_x width
+               $ flip_y height
+               $ swap_bytes arr
+  let arrSlice' = R.reshape slice3d arrSlice
+
+  -- If zoom isn't being inlined we need to extract the slice before hand,
+  -- and convert it to the require formed.
+  dbl_data   <- R.computeUnboxedP $ R.map fromIntegral arrSlice'
+  raw_output <- zoom dbl_data zoom_factor
+  arrSlice0  <- R.computeUnboxedP $ z_slice 0 raw_output
+
+  -- Make doubles from the thresholds which are given as Ints.
+  let lt = fromIntegral lower_threshold :: Double
+  let ut = fromIntegral upper_threshold :: Double
+
+  let arr_bracketed = bracket_array lt ut arrSlice0
+  values <- R.computeUnboxedP $ R.map fromIntegral arr_bracketed
+  write_values_to_bmp output values
+
+  where
+    slice3d :: R.DIM3
+    slice3d = (R.Z R.:. 1 R.:. height  R.:. width)