module PCD.Data where
import Control.Applicative
import Control.DeepSeq
import Control.Monad (when)
import Data.Attoparsec.Text
import qualified Data.Attoparsec.Text.Lazy as ATL
import qualified Data.Text.Lazy.IO as TL
import qualified Data.Text.IO as T
import qualified Data.Vector as B
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Storable as V
import qualified Data.Vector.Storable.Mutable as VM
import Foreign.Storable (Storable, sizeOf)
import System.IO (Handle, openFile, hClose,
IOMode(..), withBinaryFile, hPutBuf, hGetBuf)
import PCD.Header
import PCD.Internal.SmallLens
import PCD.Internal.Types
readAsciiPoints :: (G.Vector v a) =>
Header -> Handle -> ATL.Parser a -> IO (v a)
readAsciiPoints pcd h p = aux <$> TL.hGetContents h
where n = fromIntegral $ pcd^.points
aux t0 = G.create $
do v <- GM.new n
let write = GM.write v
go i t
| i == n = return v
| otherwise = case ATL.parse p t of
ATL.Done t' pt -> write i pt >>
go (i+1) t'
ATL.Fail _ _ msg -> error msg
go 0 t0
readAsciiPointsDefault :: Header -> Handle -> IO (B.Vector [FieldType])
readAsciiPointsDefault pcd h = readAsciiPoints pcd h $ pointParser pcd
readBinPoints :: forall a. Storable a => Header -> Handle -> IO (Either String (Vector a))
readBinPoints pcd h
| ptSize /= sz = return . Left $
"Deserialization type is not the same size as the points "++
"described by this file. The PCD file dicates "++
show ptSize++" bytes per point; destination type takes up "++
show sz++" bytes."
| otherwise = do vm <- VM.new (fromIntegral $ pcd^.points)
_ <- VM.unsafeWith vm (flip (hGetBuf h) numBytes)
Right <$> V.freeze vm
where sz = sizeOf (undefined::a)
numBytes = fromIntegral (pcd^.points) * sz
ptSize = sum (_sizes pcd)
readPointData :: Storable a =>
Header -> Handle -> ATL.Parser a ->
IO (Either String (Vector a))
readPointData hd h pa
| hd^.format == ASCII = readAsciiPoints hd h pa >>= return . Right
| otherwise = readBinPoints hd h
readXYZ_ascii :: Fractional a => ATL.Parser (V3 a)
readXYZ_ascii = (\[x,y,z] -> V3 x y z) <$>
count 3 ((realToFrac <$> double) <* skipSpace)
readXYZW_ascii :: Fractional a => ATL.Parser (V4 a)
readXYZW_ascii = (\[x,y,z,w] -> V4 x y z w) <$>
count 4 ((realToFrac <$> double) <* skipSpace)
saveBinaryPcd :: forall a. Storable a =>
FilePath -> Header -> V.Vector a -> IO ()
saveBinaryPcd outputFile pcd pts =
do putStrLn $ "Converting "++show (V.length pts)++" points"
let pcd' = format .~ Binary $ pcd
sz = sizeOf (undefined::a) * V.length pts
T.writeFile outputFile (writeHeader pcd')
withBinaryFile outputFile AppendMode $ \h ->
V.unsafeWith pts (flip (hPutBuf h) sz)
asciiToBinary :: FilePath -> FilePath -> IO ()
asciiToBinary i o = do h <- openFile i ReadMode
(pcdh,_) <- readHeader h
pcdh `deepseq` print pcdh
when ((pcdh^.format) /= ASCII)
(error "Input PCD is already binary!")
case pcdh ^. sizes of
[4,4,4] -> readAsciiPoints pcdh h readXYZ_ascii >>=
(saveBinaryPcd o pcdh
:: V.Vector (V3 Float) -> IO ())
[4,4,4,4] -> readAsciiPoints pcdh h readXYZW_ascii >>=
(saveBinaryPcd o pcdh
:: V.Vector (V4 Float) -> IO ())
_ -> error $ "Only 32-bit floating point 3D or 4D "++
"points are supported."
hClose h
loadPoints :: Storable a => ATL.Parser a -> FilePath -> IO (Vector a)
loadPoints parser pcdFile = do h <- openFile pcdFile ReadMode
(pcdh,_) <- readHeader h
r <- pcdh `deepseq` readPointData pcdh h parser
hClose h
return $ either (const V.empty) id r
loadXyz :: (Fractional a, Storable a) => FilePath -> IO (Vector (V3 a))
loadXyz = loadPoints readXYZ_ascii
loadXyzw :: (Fractional a, Storable a) => FilePath -> IO (Vector (V4 a))
loadXyzw = loadPoints readXYZW_ascii