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.1.1.1 +version: 0.2.0 synopsis: PCD file loader. description: Parser for PCD (point cloud data) formats. See @@ -18,7 +18,8 @@ location: git://github.com/acowley/pcd-loader.git library - exposed-modules: PCD.Data, PCD.Header, PCD.Internal.Types + exposed-modules: PCD.Data, PCD.Header, PCD.Internal.Types, + PCD.Internal.StorableFieldType other-modules: PCD.Internal.SmallLens ghc-options: -O2 -Wall hs-source-dirs: src diff --git a/src/PCD/Data.hs b/src/PCD/Data.hs --- a/src/PCD/Data.hs +++ b/src/PCD/Data.hs @@ -1,24 +1,30 @@ -{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE ScopedTypeVariables, BangPatterns #-} -- |Parser for PCD (point cloud data) files. Also provides a facility -- for converting from ASCII to binary formatted point data. -module PCD.Data where +module PCD.Data (FieldType(..), unsafeUnwrap, loadAllFields, + loadXyzw, loadXyz, asciiToBinary) where import Control.Applicative import Control.DeepSeq import Control.Monad (when) -import Data.Attoparsec.Text +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.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.Storable (Storable, sizeOf) +import Foreign.Marshal.Alloc (allocaBytes) +import Foreign.Ptr (Ptr, castPtr, plusPtr) +import Foreign.Storable (Storable, sizeOf, peek) 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 -- |Read point data using a user-supplied ASCII point parser. @@ -29,22 +35,24 @@ aux t0 = G.create $ do v <- GM.new n let write = GM.write v - go i t + 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.Done !t' !pt -> write i pt >> + go (i+1) t' ATL.Fail _ _ msg -> error msg go 0 t0 -- |Load points of unknown dimension into a boxed vector with a list -- of 'FieldType' as the point representation. -readAsciiPointsDefault :: Header -> Handle -> IO (B.Vector [FieldType]) -readAsciiPointsDefault pcd h = readAsciiPoints pcd h $ pointParser pcd +readAsciiPointsDefault :: Header -> Handle -> IO (B.Vector (B.Vector FieldType)) +readAsciiPointsDefault pcd h = readAsciiPoints pcd h $ + B.fromList <$> pointParser pcd -- |Read back 'Storable' points saved as binary data. -readBinPoints :: forall a. Storable a => Header -> Handle -> IO (Either String (Vector a)) -readBinPoints pcd h +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 "++ @@ -66,7 +74,7 @@ IO (Either String (Vector a)) readPointData hd h pa | hd^.format == ASCII = readAsciiPoints hd h pa >>= return . Right - | otherwise = readBinPoints hd h + | otherwise = readHomogenousBinaryPoints hd h -- |Parse 3D points serialized in ASCII. readXYZ_ascii :: Fractional a => ATL.Parser (V3 a) @@ -92,25 +100,24 @@ withBinaryFile outputFile AppendMode $ \h -> V.unsafeWith pts (flip (hPutBuf h) sz) --- |Convert the single-precision floating point XYZ or XYZW (where --- \"W\" may be an RGB triple encoded as a single float) points in an --- ASCII PCD to a binary PCD file. +-- |@asciiToBinary inputFile outputFile@ converts a PCD file from +-- ASCII to Binary. asciiToBinary :: FilePath -> FilePath -> IO () asciiToBinary i o = do h <- openFile i ReadMode (pcdh,_) <- readHeader h pcdh `deepseq` print pcdh - when ((pcdh^.format) /= ASCII) + 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." + let numBytes = totalBinarySize pcdh + putStrLn $ "Expecting to generate "++show numBytes++" bytes" + v <- readAsciiPoints pcdh h (B.fromList <$> pointParser pcdh) + 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 -- |Load points stored in a PCD file into a 'Vector'. loadPoints :: Storable a => ATL.Parser a -> FilePath -> IO (Vector a) @@ -134,3 +141,21 @@ loadXyzw = loadPoints readXYZW_ascii {-# SPECIALIZE loadXyzw :: FilePath -> IO (Vector (V4 Float)) #-} {-# SPECIALIZE loadXyzw :: FilePath -> IO (Vector (V4 Double)) #-} + +-- |Parse every field of every point in a PCD file. Returns a function +-- that may be used to project out a named field. +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) diff --git a/src/PCD/Header.hs b/src/PCD/Header.hs --- a/src/PCD/Header.hs +++ b/src/PCD/Header.hs @@ -27,15 +27,51 @@ -- |The PCD format has both ASCII and binary variants. data DataFormat = ASCII | Binary deriving (Eq,Show,Ord) -data FieldType = TUchar Word8 - | TChar Int8 - | TUshort Word16 - | TShort Int16 - | TUint Word32 - | TInt Int32 - | TFloat Float - | TDouble Double +data FieldType = TUchar {-# UNPACK #-}!Word8 + | TChar {-# UNPACK #-}!Int8 + | TUshort {-# UNPACK #-}!Word16 + | TShort {-# UNPACK #-}!Int16 + | TUint {-# UNPACK #-}!Word32 + | TInt {-# UNPACK #-}!Int32 + | TFloat {-# UNPACK #-}!Float + | TDouble {-# UNPACK #-}!Double + deriving Show +class PCDType a where + unsafeUnwrap :: FieldType -> a + +instance PCDType Word8 where + unsafeUnwrap (TUchar x) = x + unsafeUnwrap y = error $ "Tried to unwrap "++show y++" as a Word8" + +instance PCDType Int8 where + unsafeUnwrap (TChar x) = x + unsafeUnwrap y = error $ "Tried to unwrap "++show y++" as a Int8" + +instance PCDType Word16 where + unsafeUnwrap (TUshort x) = x + unsafeUnwrap y = error $ "Tried to unwrap "++show y++" as a Word16" + +instance PCDType Int16 where + unsafeUnwrap (TShort x) = x + unsafeUnwrap y = error $ "Tried to unwrap "++show y++" as a Int16" + +instance PCDType Word32 where + unsafeUnwrap (TUint x) = x + unsafeUnwrap y = error $ "Tried to unwrap "++show y++" as a Word32" + +instance PCDType Int32 where + unsafeUnwrap (TInt x) = x + unsafeUnwrap y = error $ "Tried to unwrap "++show y++" as a Int32" + +instance PCDType Float where + unsafeUnwrap (TFloat x) = x + unsafeUnwrap y = error $ "Tried to unwrap "++show y++" as a Float" + +instance PCDType Double where + unsafeUnwrap (TDouble x) = x + unsafeUnwrap y = error $ "Tried to unwrap "++show y++" as a Double" + -- |Construct a parser for a field based on its type and size. fieldParser :: DimType -> Int -> Parser FieldType fieldParser I 1 = TChar <$> (decimal :: Parser Int8) @@ -48,10 +84,18 @@ fieldParser F 8 = TDouble <$> double fieldParser t s = error $ "Unknown field description: "++show t++" "++show s +sequence' :: Monad m => [m a] -> m [a] +sequence' [] = return [] +sequence' (x:xs) = do !x' <- x + !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 $ zipWith fieldParser (_dimTypes h) (_sizes h) +pointParser h = sequence' . map (<* skipSpace) $ + zipWith fieldParser (_dimTypes h) (_sizes h) data Header = Header { _version :: Text , _fields :: [Text] @@ -147,3 +191,9 @@ tshow = T.pack . show ftshow :: (Foldable t, Show a) => t a -> Text ftshow = mconcat . intersperse " " . foldMap ((:[]) . tshow) + +-- |Compute the number of bytes this point cloud would occupy if +-- serialized with the 'Binary' encoding. +totalBinarySize :: Header -> Int +totalBinarySize h = fromIntegral (h^.points) * + sum (zipWith (*) (h^.counts) (h^.sizes)) diff --git a/src/PCD/Internal/StorableFieldType.hs b/src/PCD/Internal/StorableFieldType.hs new file mode 100644 --- /dev/null +++ b/src/PCD/Internal/StorableFieldType.hs @@ -0,0 +1,85 @@ +{-# LANGUAGE ScopedTypeVariables, BangPatterns #-} +module PCD.Internal.StorableFieldType where +import Control.Applicative +import Control.Lens ((^.)) +import Control.Monad (void) +import qualified Data.Vector as B +import qualified Data.Vector.Mutable as BM +import PCD.Header +import Foreign.Marshal.Alloc (allocaBytes) +import Foreign.Ptr (Ptr, castPtr, plusPtr) +import Foreign.Storable (Storable, peek, poke, sizeOf) +import System.IO (Handle, hGetBuf) + +data P a = P !FieldType {-# UNPACK #-} !(Ptr a) + +-- |'peek' a 'Storable' and advance the source pointer past this +-- datum. +peekStep :: forall a b. Storable a => (a -> FieldType) -> Ptr b -> IO (P b) +peekStep mk ptr = P . mk <$> peek (castPtr ptr) + <*> pure (plusPtr ptr (sizeOf (undefined::a))) + +parseBinaryField :: DimType -> Int -> Ptr a -> IO (P a) +parseBinaryField I 1 = peekStep TChar +parseBinaryField I 2 = peekStep TShort +parseBinaryField I 4 = peekStep TInt +parseBinaryField U 1 = peekStep TUchar +parseBinaryField U 2 = peekStep TUshort +parseBinaryField U 4 = peekStep TUint +parseBinaryField F 4 = peekStep TFloat +parseBinaryField F 8 = peekStep TDouble + +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 + where n = fromIntegral $ pcdh ^. points + numBytes = n * sum (zipWith (*) (pcdh^.counts) (pcdh^.sizes)) + pointParser = parseBinaryFields pcdh + +parseBinaryFields :: Header -> Ptr a -> IO (B.Vector FieldType, Ptr a) +parseBinaryFields h ptr = aux ptr + where numFields = sum (h^.counts) + aux ptr0 = (\(v,ptr) -> (,) <$> B.unsafeFreeze v <*> pure ptr) =<< + do v <- BM.new numFields + let write = BM.write v + go !i !ptr ss ts cs + | i == numFields = return (v,ptr) + | otherwise = + do P x ptr' <- parseBinaryField (head ts) + (head ss) + ptr + write i x + let (c:cs') = cs + if c == 1 + then go (i+1) ptr' (tail ss) (tail ts) cs' + else go (i+1) ptr' ss ts (c-1 : cs') + go 0 ptr0 (h^.sizes) (h^.dimTypes) (h^.counts) + +pokeStep :: forall a b. Storable a => a -> Ptr b -> IO (Ptr b) +pokeStep x ptr = poke (castPtr ptr) x >> + return (plusPtr ptr (sizeOf (undefined::a))) + +pokeBinaryField :: FieldType -> Ptr a -> IO (Ptr a) +pokeBinaryField (TUchar x) = pokeStep x +pokeBinaryField (TChar x) = pokeStep x +pokeBinaryField (TUshort x) = pokeStep x +pokeBinaryField (TShort x) = pokeStep x +pokeBinaryField (TUint x) = pokeStep x +pokeBinaryField (TInt x) = pokeStep x +pokeBinaryField (TFloat x) = pokeStep x +pokeBinaryField (TDouble x) = pokeStep x + +pokeBinaryFields :: Ptr a -> B.Vector FieldType -> IO (Ptr a) +pokeBinaryFields = B.foldM' aux + where aux ptr x = pokeBinaryField x ptr + +pokeBinaryPoints :: Ptr a -> B.Vector (B.Vector FieldType) -> IO () +pokeBinaryPoints = (void .) . B.foldM' pokeBinaryFields