module PCD.Data (FieldType(..), unsafeUnwrap, loadAllFields,
loadXyzw, loadXyz, asciiToBinary, saveBinaryPcd,
projectBinaryFields, mkSimpleHeader, mkHeaderXYZ) where
import Control.Applicative
import Control.DeepSeq
import Control.Monad (when)
import Data.Attoparsec.Text hiding (I)
import qualified Data.Attoparsec.Text.Lazy as ATL
import Data.Text (Text)
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.Marshal.Alloc (allocaBytes)
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.StorableFieldType
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 (B.Vector FieldType))
readAsciiPointsDefault pcd h = readAsciiPoints pcd h $
B.fromList <$> pointParser pcd
readHomogenousBinaryPoints :: forall a. Storable a =>
Header -> Handle -> IO (Either String (Vector a))
readHomogenousBinaryPoints 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 = readHomogenousBinaryPoints 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!")
let numBytes = totalBinarySize pcdh
putStrLn $ "Expecting to generate "++show numBytes++" bytes"
v <- readAsciiPointsDefault pcdh h
putStrLn $ "Parsed "++show (B.length v)++" ASCII points"
hClose h
T.writeFile o (writeHeader (format .~ Binary $ pcdh))
withBinaryFile o AppendMode $ \h' ->
allocaBytes numBytes $ \ptr ->
pokeBinaryPoints ptr v >>
hPutBuf h' ptr numBytes
projectBinaryFields :: [Text] -> FilePath -> FilePath -> IO ()
projectBinaryFields fs i o =
do h <- openFile i ReadMode
(pcdh,_) <- readHeader h
v <- loadFlexiblePoints pcdh h
putStrLn $ "Parsed "++show (B.length v)++" ASCII points"
let v' = B.map keep v
keepers = B.fromList $ map (`elem` fs) (pcdh ^. fields)
keep = B.map snd . B.filter fst . B.zip keepers
pcdh' = format .~ Binary $ filterFields (`elem` fs) pcdh
numBytes = totalBinarySize pcdh'
putStrLn $ "Binary data will occupy "++show numBytes++" bytes"
hClose h
T.writeFile o (writeHeader pcdh')
withBinaryFile o AppendMode $ \h' ->
allocaBytes numBytes $ \ptr ->
pokeBinaryPoints ptr v' >>
hPutBuf h' ptr numBytes
return ()
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
loadFlexiblePoints :: Header -> Handle -> IO (B.Vector (B.Vector FieldType))
loadFlexiblePoints pcdh h
| pcdh ^. format == Binary = parseBinaryPoints pcdh h
| otherwise = readAsciiPointsDefault pcdh h
loadAllFields :: FilePath -> IO (Text -> B.Vector FieldType)
loadAllFields f = do h <- openFile f ReadMode
(pcdh,_) <- readHeader h
(mkProjector pcdh <$>
if pcdh ^. format == ASCII
then readAsciiPoints pcdh h
(B.fromList <$> pointParser pcdh)
else parseBinaryPoints pcdh h)
<* hClose h
where mkProjector :: Header -> B.Vector (B.Vector FieldType) ->
(Text -> B.Vector FieldType)
mkProjector h pts = let fieldNames = B.fromList $ h ^. fields
in \name -> maybe B.empty
(flip B.map pts . flip (B.!))
(B.findIndex (name==) fieldNames)