module PCD.Data (FieldType(..), unsafeUnwrap, loadAllFields,
loadXyzw, loadXyz, asciiToBinary, saveBinaryPcd,
projectBinaryFields, mkSimpleHeader, mkHeaderXYZ) where
import Control.Applicative
import Control.DeepSeq
import Control.Lens ((.~), (^.))
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.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)