]> gitweb.michael.orlitzky.com - spline3.git/blob - src/Main.hs
5e10d6cfdc16146d7d8844de335558253c6ef5b2
[spline3.git] / 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 word16_output <- R.computeUnboxedP $ round_array raw_output
93 write_word16s output word16_output
94
95
96 main2d :: Args -> R.DIM3 -> IO ()
97 main2d Args{..} shape = do
98 let zoom_factor = (1, scale, scale)
99 arr <- read_word16s input shape
100 arrSlice <- R.computeUnboxedP
101 $ z_slice (fromJust slice)
102 $ flip_x width
103 $ flip_y height
104 $ swap_bytes arr
105 let arrSlice' = R.reshape slice3d arrSlice
106
107 -- If zoom isn't being inlined we need to extract the slice before hand,
108 -- and convert it to the require formed.
109 dbl_data <- R.computeUnboxedP $ R.map fromIntegral arrSlice'
110 raw_output <- zoom dbl_data zoom_factor
111 arrSlice0 <- R.computeUnboxedP $ z_slice 0 raw_output
112
113 -- Make doubles from the thresholds which are given as Ints.
114 let lt = fromIntegral lower_threshold
115 let ut = fromIntegral upper_threshold
116
117 let arr_bracketed = bracket_array lt ut arrSlice0
118 values <- R.computeUnboxedP $ R.map fromIntegral arr_bracketed
119 write_values_to_bmp output values
120
121 where
122 slice3d :: R.DIM3
123 slice3d = (R.Z R.:. 1 R.:. height R.:. width)