diff --git a/pcd-loader.cabal b/pcd-loader.cabal --- a/pcd-loader.cabal +++ b/pcd-loader.cabal @@ -1,5 +1,5 @@ name: pcd-loader -version: 0.2.0 +version: 0.2.2 synopsis: PCD file loader. description: Parser for PCD (point cloud data) formats. See diff --git a/src/PCD/Data.hs b/src/PCD/Data.hs --- a/src/PCD/Data.hs +++ b/src/PCD/Data.hs @@ -2,7 +2,8 @@ -- |Parser for PCD (point cloud data) files. Also provides a facility -- for converting from ASCII to binary formatted point data. module PCD.Data (FieldType(..), unsafeUnwrap, loadAllFields, - loadXyzw, loadXyz, asciiToBinary) where + loadXyzw, loadXyz, asciiToBinary, saveBinaryPcd, + projectBinaryFields) where import Control.Applicative import Control.DeepSeq import Control.Monad (when) @@ -12,14 +13,12 @@ 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.Mutable as BM 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.Ptr (Ptr, castPtr, plusPtr) -import Foreign.Storable (Storable, sizeOf, peek) +import Foreign.Storable (Storable, sizeOf) import System.IO (Handle, openFile, hClose, IOMode(..), withBinaryFile, hPutBuf, hGetBuf) import PCD.Header @@ -110,7 +109,7 @@ (error "Input PCD is already binary!") let numBytes = totalBinarySize pcdh putStrLn $ "Expecting to generate "++show numBytes++" bytes" - v <- readAsciiPoints pcdh h (B.fromList <$> pointParser pcdh) + v <- readAsciiPointsDefault pcdh h putStrLn $ "Parsed "++show (B.length v)++" ASCII points" hClose h T.writeFile o (writeHeader (format .~ Binary $ pcdh)) @@ -119,7 +118,31 @@ pokeBinaryPoints ptr v >> hPutBuf h' ptr numBytes --- |Load points stored in a PCD file into a 'Vector'. +-- |Save a binary PCD file including only the named fields. +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 () + +-- |Load points stored in a PCD file into a 'Vector'. This requires a +-- 'Storable' instance for the type used to represent a point. If the +-- point is a monotyped collection of fields, consider using +-- 'Linear.V2.V2' or 'Linear.V3.V3' to represent points. loadPoints :: Storable a => ATL.Parser a -> FilePath -> IO (Vector a) loadPoints parser pcdFile = do h <- openFile pcdFile ReadMode (pcdh,_) <- readHeader h @@ -141,6 +164,11 @@ loadXyzw = loadPoints readXYZW_ascii {-# SPECIALIZE loadXyzw :: FilePath -> IO (Vector (V4 Float)) #-} {-# SPECIALIZE loadXyzw :: FilePath -> IO (Vector (V4 Double)) #-} + +loadFlexiblePoints :: Header -> Handle -> IO (B.Vector (B.Vector FieldType)) +loadFlexiblePoints pcdh h + | pcdh ^. format == Binary = parseBinaryPoints pcdh h + | otherwise = readAsciiPointsDefault pcdh h -- |Parse every field of every point in a PCD file. Returns a function -- that may be used to project out a named field. diff --git a/src/PCD/Header.hs b/src/PCD/Header.hs --- a/src/PCD/Header.hs +++ b/src/PCD/Header.hs @@ -90,13 +90,6 @@ !xs' <- sequence' xs return $ x':xs' - --- |Assemble a parser for points by sequencing together all necessary --- field parsers. -pointParser :: Header -> Parser [FieldType] -pointParser h = sequence' . map (<* skipSpace) $ - zipWith fieldParser (_dimTypes h) (_sizes h) - data Header = Header { _version :: Text , _fields :: [Text] , _sizes :: [Int] @@ -108,6 +101,20 @@ , _points :: Integer , _format :: DataFormat } deriving Show makeLenses ''Header + +-- |Assemble a parser for points by sequencing together all necessary +-- field parsers. +pointParser :: Header -> Parser [FieldType] +pointParser h = sequence' . map (<* skipSpace) $ + zipWith fieldParser (_dimTypes h) (_sizes h) + +-- |Create a 'Header' based on an existing one that keeps only the +-- fields whose names pass the supplied predicate. +filterFields :: (Text -> Bool) -> Header -> Header +filterFields f h = h % (fields %~ keep) . (sizes %~ keep) + . (dimTypes %~ keep) . (counts %~ keep) + where keepers = map f (_fields h) + keep = map snd . filter fst . zip keepers instance NFData Header where rnf (Header !_v !_f !_s !_d !_c !_w !_h !(!_t,!_r) !_p !_fmt) = () diff --git a/src/PCD/Internal/StorableFieldType.hs b/src/PCD/Internal/StorableFieldType.hs --- a/src/PCD/Internal/StorableFieldType.hs +++ b/src/PCD/Internal/StorableFieldType.hs @@ -1,5 +1,5 @@ {-# LANGUAGE ScopedTypeVariables, BangPatterns #-} -module PCD.Internal.StorableFieldType where +module PCD.Internal.StorableFieldType (parseBinaryPoints, pokeBinaryPoints) where import Control.Applicative import Control.Lens ((^.)) import Control.Monad (void) @@ -11,6 +11,7 @@ import Foreign.Storable (Storable, peek, poke, sizeOf) import System.IO (Handle, hGetBuf) +-- Strict tuple used during parsing. data P a = P !FieldType {-# UNPACK #-} !(Ptr a) -- |'peek' a 'Storable' and advance the source pointer past this @@ -28,24 +29,29 @@ parseBinaryField U 4 = peekStep TUint parseBinaryField F 4 = peekStep TFloat parseBinaryField F 8 = peekStep TDouble +parseBinaryField t s = error ("Unknown field type: "++show t++" "++show s) parseBinaryPoints :: Header -> Handle -> IO (B.Vector (B.Vector FieldType)) -parseBinaryPoints pcdh h = B.unsafeFreeze =<< - do v <- BM.new n - let go !i !ptr - | i == n = return v - | otherwise = do (pt,ptr') <- pointParser ptr - BM.write v i pt - go (i+1) ptr' - allocaBytes numBytes $ \ptr -> - hGetBuf h ptr numBytes >> - go 0 ptr +parseBinaryPoints pcdh h = + B.unsafeFreeze =<< + do v <- BM.new n + let go !i !ptr + | i == n = return v + | otherwise = do (pt,ptr') <- pointParserBin ptr + BM.write v i pt + go (i+1) ptr' + allocaBytes numBytes $ \ptr -> + hGetBuf h ptr numBytes >> + go 0 ptr where n = fromIntegral $ pcdh ^. points numBytes = n * sum (zipWith (*) (pcdh^.counts) (pcdh^.sizes)) - pointParser = parseBinaryFields pcdh + pointParserBin = parseBinaryFields pcdh +-- Parse all fields of a single point. A point is represented as a +-- 'B.Vector' of its fields. The returned 'Ptr' is just after the +-- parsed point. parseBinaryFields :: Header -> Ptr a -> IO (B.Vector FieldType, Ptr a) -parseBinaryFields h ptr = aux ptr +parseBinaryFields h = aux where numFields = sum (h^.counts) aux ptr0 = (\(v,ptr) -> (,) <$> B.unsafeFreeze v <*> pure ptr) =<< do v <- BM.new numFields diff --git a/src/executable/Main.hs b/src/executable/Main.hs --- a/src/executable/Main.hs +++ b/src/executable/Main.hs @@ -1,10 +1,22 @@ +{-# LANGUAGE OverloadedStrings #-} module Main (main) where -import Control.Monad (when) import System.Environment (getArgs) -import PCD.Data (asciiToBinary) +import PCD.Data (asciiToBinary, projectBinaryFields) +data Args = Args { _inputFile :: FilePath + , _outputFile :: FilePath + , _justXyz :: Bool } + +parseArgs :: [String] -> Maybe Args +parseArgs ["-xyz", a, b] = Just $ Args a b True +parseArgs [a,b] = Just $ Args a b False +parseArgs _ = Nothing + +usage :: IO () +usage = do putStrLn "Usage: pcd2bin [-xyz] asciiPcd outputBinaryFile" + putStrLn "- The '-xyz' option restricts output to those fields." + main :: IO () -main = do args@(~[inputFile, outputFile]) <- getArgs - when (length args /= 2) - (error "Usage: pcd2bin asciiPcd outputBinaryFile") - asciiToBinary inputFile outputFile +main = getArgs >>= maybe usage aux . parseArgs + where aux (Args a b False) = asciiToBinary a b + aux (Args a b True) = projectBinaryFields ["x","y","z"] a b