explore-0.0.5.1: Main.hs

-- (c) 2009, Cetin Sert
-- | ExPloRe: Experimental Plot Reconstructor

module Main where

import System.Environment
import System.IO.Unsafe
import System.Exit
import Data.Word
import Foreign.Ptr
import Foreign.Storable
import Data.Array.Storable
import Control.Monad
import Codec.Image.PNG

type Name = String

-- | 32-bit color representation
data RGBA = RGBA Word8 Word8 Word8 Word8 -- ^ RED BLUE GREEN ALPHA
          deriving (Show, Read, Eq)

-- storable instance for retrieveal 
instance Storable RGBA where
  sizeOf _    = sizeOf (0 :: Word8) * 4
  alignment _ = 1
  peek color  = do
	let byte :: Ptr Word8 = castPtr color
	[r,g,b,a] <- mapM (byte @!) [0..3]
	return $ RGBA r g b a

-- | colour matching behaviour
data Matching = S                  -- ^ strict
              | TA Int             -- ^ shared threshold for all color components
              | TC Int Int Int Int -- ^ individual thresholds for color components
              deriving (Show, Read, Eq)

-- | scan area
--   w = width of the PNG image
--   h = height of the PNG image
data ScanArea = F                  -- ^ search all the area (FULL)        | 0 0 w h
              | O Int Int          -- ^ offset from LEFT TOP              | L T w h
              | M Int Int Int Int  -- ^ offset from LEFT TOP RIGHT BOTTOM | L T (w-R) (h-B)
              | A Int Int Int Int  -- ^ area        LEFT TOP RIGHT BOTTOM | L T R B
              deriving (Show, Read, Eq)

-- | normalizes scan area representations
normalize :: Int -> Int -> ScanArea -> ScanArea
normalize w h F           = A 0 0 w h
normalize w h (O l t)     = A l t w h
normalize w h (M l t r b) = A l t (w-r) (h-b)
normalize _ _ area        = area

-- | ExPloRe requires PNGs with transparency channel and no interlacing 
checkAlpha :: PNGImage -> IO ()
checkAlpha img = case hasAlphaChannel img of
  True -> return ()
  _    -> putStrLn "no alpha channel going on!!!" >> exitWith (ExitFailure 1)


main :: IO ()
main = do
  putStrLn $ "ExPloRe 0.0.5.1 : Experimental Plot Reconstructor"
  args@(imgPath:legendPath:matching_:area_:step_:_) <- getArgs

  -- initialize image
  Right img <- loadPNGFile imgPath
  let bitmap  = imageData  img
  let (wu,hu) = dimensions img
  let (w,h)   = (fromIntegral wu, fromIntegral hu)
  putStrLn $ "-------------------------------------------------------------------"
  putStrLn $ ""
  putStrLn $ "call  : " ++ tail (filter (/= '"') $ concatMap ((' ':) . show) args)
  putStrLn $ ""
  putStrLn $ "image : " ++ imgPath
  putStrLn $ "width : " ++ show w
  putStrLn $ "height: " ++ show h
  putStrLn $ ""
  checkAlpha img 

  -- initialize scan
  let area@(A l t r b) = normalize w h . read $ area_ :: ScanArea
  let start = fromIntegral l :: Double
  let step  = read step_ :: Double
  let (@#)  = make2DIndexer w
  let rows  = [t..b]
  let cols  = takeWhile (< r) $ map (floor . (start +) . (step *)) [0..]
  let icols = zip [1..] cols
  let matching = read matching_ :: Matching
  let (~=)  = colorEq matching
  putStrLn $ "area  : " ++ show area
  putStrLn $ "match : " ++ show matching

  -- initialize lines
  lines_ <- readFile legendPath
  let lines = read lines_ :: [(Name,RGBA)]
  putStrLn $ "legend: " ++ legendPath
  putStrLn $ "lines : " ++ (show $ length lines)
  putStrLn $ "step  : " ++ show step
  putStrLn $ ""
  mapM_ (putStrLn . show) lines

  -- scan bitmap
  mapM_ (scan bitmap area icols rows (@#) (~=)) lines


-- | a very simple scanning algorithm that works only if no lines
--   occlude each other.
scan bitmap (A l t r b) icols rows (@#) (~=) (name,color) = do
  putStrLn $ ""
  putStrLn $ "-------------------------------------------------------------------"
  putStrLn $ show color
  putStrLn $ ""
  putStrLn $ name
  putStrLn $ ""
  putStrLn $ "point\tx\tmid y\tmatches y"
  putStrLn $ ""
  withStorableArray bitmap $ \byte -> do
    let pixel :: Ptr RGBA = castPtr byte
    forM_ icols $ \(n,j) -> do
      let matches = flip filter rows $ \i -> (pixel @# i) j ~= color
      let m = mid matches - t
      putStrLn $ case not . null $ matches of
        True  -> show n ++ "\t" ++ show j ++ "\t" ++ show m ++ "\t" ++ show matches
        False -> show n ++ "\t" ++ show j ++ "\t   \t[]"


-- | fuzzy color matching
colorEq :: Matching -> RGBA -> RGBA -> Bool
colorEq S            = compEq 0 0 0 0
colorEq (TA v)       = compEq v v v v
colorEq (TC r g b a) = compEq r g b a

-- | fuzzy color matching
compEq :: Int -> Int -> Int -> Int -> RGBA -> RGBA -> Bool
compEq tr tg tb ta (RGBA a b c d) (RGBA x y z w) = let fe = fuzzyEq in
  fe tr a x && fe tg b y && fe tb c z && fe ta d w

-- | simple fuzzy equality
fuzzyEq :: forall b a. (Num b, Ord b, Integral a) => b -> a -> a -> Bool
fuzzyEq 0 = (==)
fuzzyEq t = \x y -> t > (abs $ fromIntegral x - fromIntegral y)

-- | gets the middle element of a list
mid :: forall a. [a] -> a
mid xs = xs !! (fromIntegral . floor . (/ 2) . fromIntegral . length) xs

(@!) :: Storable a => Ptr a -> Int -> IO a
(@!) = peekElemOff

make2DIndexer :: Storable a => Int -> Ptr a -> Int -> Int -> a
make2DIndexer w = \p j i -> unsafePerformIO $ p @! (i + j * w)