module PCD.Data (
FieldType(..), unsafeUnwrap,
loadFieldsByName, loadFlexiblePoints, loadXyzw, loadXyz,
saveBinaryPcd, projectBinaryFields,
mkSimpleHeader, mkHeaderXYZ) where
import Control.Applicative
import Control.DeepSeq
import Control.Lens ((.~), (^.))
import qualified Data.Attoparsec.Text.Lazy as ATL
import Data.Text (Text)
import qualified Data.Text.IO as T
import qualified Data.Vector as B
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 qualified PCD.Internal.AsciiParsers as A
import PCD.Internal.StorableFieldType
import PCD.Internal.Types
readStorableBinaryPoints :: forall a. Storable a =>
Header -> Handle -> IO (Either String (Vector a))
readStorableBinaryPoints 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 header handle parser
| header^.format == ASCII = Right <$> A.readPoints header handle parser
| otherwise = readStorableBinaryPoints header handle
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)
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
v <- case r of
Left _ -> return V.empty
Right v' -> V.length v' `seq` return v'
hClose h
return v
loadXyz :: (Fractional a, Storable a) => FilePath -> IO (Vector (V3 a))
loadXyz = loadPoints A.readXYZ
loadXyzw :: (Fractional a, Storable a) => FilePath -> IO (Vector (V4 a))
loadXyzw = loadPoints A.readXYZW
loadFlexiblePoints :: Header -> Handle -> IO (B.Vector (B.Vector FieldType))
loadFlexiblePoints pcdh h
| pcdh ^. format == Binary = parseBinaryPoints pcdh h
| otherwise = A.readPointsDefault pcdh h
loadFieldsByName :: FilePath -> IO (Text -> B.Vector FieldType)
loadFieldsByName f = do h <- openFile f ReadMode
(pcdh,_) <- readHeader h
(mkProjector pcdh <$> loadFlexiblePoints 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)