-- GENERATED by C->Haskell Compiler, version 0.16.3 Crystal Seed, 24 Jan 2009 (Haskell)
-- Edit the ORIGNAL .chs file instead!


{-# LINE 1 "./CV/Calibration.chs" #-}{-#LANGUAGE ForeignFunctionInterface, ScopedTypeVariables #-}
-- | This module exposes opencv functions for camera calibration using a chessboard rig. This module follows opencv quite closely and the best documentation
--   is probably found there. As quick example however, the following program detects and draws chessboard corners from an image.
--
-- @
-- module Main where
-- import CV.Image
-- import CV.Calibration
-- 
-- main = do
--     Just i <- loadColorImage \"chess.png\"
--     let corners = findChessboardCorners (unsafeImageTo8Bit i) (4,5) (FastCheck:defaultFlags)
--     let y = drawChessboardCorners (unsafeImageTo8Bit i) (4,5) corners
--     mapM_ print (corners)
--     saveImage \"found_chessboard.png\" y
-- @
module CV.Calibration 
    (
     -- * Finding chessboard calibration rig
     FindFlags(..)
    ,defaultFlags
    ,findChessboardCorners
    ,refineChessboardCorners
    -- * Visualization
    ,drawChessboardCorners
    -- * Camera calibration
    ,calibrateCamera2) where
{-#OPTIONS-GHC -fwarn-unused-imports #-}

import Foreign.C.Types
import Foreign.C.String
import Foreign.ForeignPtr
import Foreign.Storable
import Foreign.Marshal.Array
import Foreign.Marshal.Utils
import Foreign.Ptr
import Data.Bits

import CV.Image 

import System.IO.Unsafe
import Utils.Point
import Control.Applicative

import CV.Matrix
import CV.Bindings.Calibrate
import CV.Bindings.Types

import CV.Image
{-# LINE 50 "./CV/Calibration.chs" #-}

-- | Flags for the chessboard corner detector. See opencv documentation for cvFindChessboardCorners.
data FindFlags = AdaptiveThresh
               | NormalizeImage
               | FilterQuads
               | FastCheck
               
instance Enum FindFlags where
  fromEnum AdaptiveThresh = 1
  fromEnum NormalizeImage = 2
  fromEnum FilterQuads = 4
  fromEnum FastCheck = 8

  toEnum 1 = AdaptiveThresh
  toEnum 2 = NormalizeImage
  toEnum 4 = FilterQuads
  toEnum 8 = FastCheck
  toEnum unmatched = error ("FindFlags.toEnum: Cannot match " ++ show unmatched)

{-# LINE 62 "./CV/Calibration.chs" #-}

flagsToNum fs = foldl (.|.) 0 $ map (fromIntegral . fromEnum) fs

-- |Default flags for finding corners
defaultFlags :: [FindFlags]
defaultFlags = [AdaptiveThresh]

-- | Find the inner corners of a chessboard in a given image. 
findChessboardCorners :: CV.Image.Image RGB D8 -> (Int, Int) -> [FindFlags] -> [(Float,Float)]
findChessboardCorners image (w,h) flags =
   unsafePerformIO $ 
    with 1 $ \(c_corner_count::Ptr CInt) -> 
     allocaArray len $ \(c_corners :: Ptr CvPoint )-> 
      withGenImage image $ \c_image -> do
        r <- wrapFindChessBoardCorners c_image (fromIntegral w) (fromIntegral h)
                                           (castPtr c_corners) c_corner_count 
                                           (flagsToNum flags)
        count <- peek c_corner_count
        arr <- peekArray (fromIntegral count) c_corners
        return (map cvPt2Pt arr) 
  where len = w*h

-- |Given an estimate of chessboard corners, provide a subpixel estimation of actual corners.
refineChessboardCorners :: Image GrayScale D8 -> [(Float,Float)] -> (Int,Int) -> (Int,Int) -> [(Float,Float)]
refineChessboardCorners img pts (winW,winH) (zeroW,zeroH) = unsafePerformIO $ do
    with 1 $ \(c_corner_count::Ptr CInt) -> 
      withImage img $ \c_img ->
      withArray (map mkPt pts) $ \(c_corners :: Ptr C'CvPoint2D32f ) -> do 
        c'wrapFindCornerSubPix c_img c_corners (length pts) winW winH zeroW zeroH tType maxIter epsilon 
        map fromPt `fmap` peekArray (length pts) c_corners
 where
    mkPt (x,y) = C'CvPoint2D32f (realToFrac x) (realToFrac y)
    fromPt (C'CvPoint2D32f x y) = (realToFrac x,realToFrac y)
    tType = c'CV_TERMCRIT_ITER
    maxIter = 100
    epsilon = 0.01

-- | Draw the found chessboard corners to an image
drawChessboardCorners :: CV.Image.Image RGB D8 -> (Int, Int) -> [(Float,Float)] -> CV.Image.Image RGB D8
drawChessboardCorners image (w,h) corners =
   unsafePerformIO $ 
    withCloneValue image $ \clone -> 
     withArray (map pt2CvPt corners) $ \(c_corners :: Ptr CvPoint )-> 
      withGenImage clone$ \c_image -> do
        r <- wrapDrawChessBoardCorners c_image (fromIntegral w) (fromIntegral h)
                                           (castPtr c_corners) (fromIntegral $ length corners) 
                                           (found)
        return clone
  where 
    len = w*h
    found |(w*h) == length corners = 1
          | otherwise = 0 
    
newtype CvPoint = CvPt (CFloat,CFloat) deriving (Show)
cvPt2Pt (CvPt (a,b)) = (realToFrac a , realToFrac b)
pt2CvPt (a,b) = CvPt (realToFrac a , realToFrac b)

instance Storable CvPoint where
  sizeOf _ = 8
{-# LINE 121 "./CV/Calibration.chs" #-}
  alignment _ = 4
{-# LINE 122 "./CV/Calibration.chs" #-}
  peek p = CvPt <$> ((,) 
    <$> (\ptr -> do {peekByteOff ptr 0 ::IO CFloat}) p
    <*> (\ptr -> do {peekByteOff ptr 4 ::IO CFloat}) p)
  poke p (CvPt (hx,hy)) = do
    (\ptr val -> do {pokeByteOff ptr 0 (val::CFloat)}) p (hx)
    (\ptr val -> do {pokeByteOff ptr 4 (val::CFloat)}) p (hy)

-- | See opencv function cvCalibrateCamera2. This function takes a list of world-screen coordinate pairs acquired with find-chessboard corners
--   and attempts to find the camera parameters for the system. It returns the fitting error, the camera matrix, list of distortion co-efficients
--   and rotation and translation vectors for each coordinate pair. 
calibrateCamera2 ::
     [[((Float, Float, Float), (Float, Float))]]
     -> (Int, Int)
     -> IO (Double, Matrix Float, [[Float]], [[Float]], [[Float]])
calibrateCamera2 views (w,h) = do
    let 
        pointCounts :: Matrix Int
        pointCounts  = fromList (1,length views) (map (length) views)
        m = length views
        totalPts = length (concat views)
        objectPoints :: Matrix Float
        objectPoints = fromList (3,totalPts) $ concat [[x,y,z] | ((x,y,z),_) <- concat views]
        imagePoints :: Matrix Float
        imagePoints  = fromList (2,totalPts) $ concat [[x,y]   | (_,(x,y))   <- concat views]
        flags = c'CV_CALIB_FIX_K1
                .|.  c'CV_CALIB_FIX_K1
                .|.  c'CV_CALIB_FIX_K2
                .|.  c'CV_CALIB_FIX_K3
                .|.  c'CV_CALIB_FIX_K4
                .|.  c'CV_CALIB_FIX_K5
                .|.  c'CV_CALIB_FIX_K6
                .|.  c'CV_CALIB_ZERO_TANGENT_DIST

        size = C'CvSize (fromIntegral w) (fromIntegral h)
        cameraMatrix,distCoeffs,rvecs,tvecs :: Matrix Float
        cameraMatrix = emptyMatrix (3,3)
        distCoeffs   = emptyMatrix (1,8)
        rvecs        = emptyMatrix (m,3)
        tvecs        = emptyMatrix (m,3)

    err <- with size $\c_size ->
     withMatPtr objectPoints $ \c_objectPoints ->
     withMatPtr imagePoints $ \c_imagePoints ->
     withMatPtr pointCounts $ \c_pointCounts ->
     withMatPtr cameraMatrix $ \c_cameraMatrix ->
     withMatPtr distCoeffs $ \c_distCoeffs ->
     withMatPtr rvecs $ \c_rvecs ->
     withMatPtr tvecs $ \c_tvecs ->
      c'wrapCalibrateCamera2 c_objectPoints c_imagePoints c_pointCounts c_size 
                             c_cameraMatrix c_distCoeffs c_rvecs c_tvecs flags

    -- print ( objectPoints, imagePoints, pointCounts,cameraMatrix, distCoeffs, rvecs, tvecs )
    return (err, transpose cameraMatrix, toCols distCoeffs, toCols rvecs, toCols tvecs)


foreign import ccall safe "CV/Calibration.chs.h wrapFindChessBoardCorners"
  wrapFindChessBoardCorners :: ((Ptr ()) -> (CInt -> (CInt -> ((Ptr ()) -> ((Ptr CInt) -> (CInt -> (IO CInt)))))))

foreign import ccall safe "CV/Calibration.chs.h wrapDrawChessBoardCorners"
  wrapDrawChessBoardCorners :: ((Ptr ()) -> (CInt -> (CInt -> ((Ptr ()) -> (CInt -> (CInt -> (IO CInt)))))))