modularized also ASM IO

unibas-gravis · Feb 2, 2024 · 8b51027 · 8b51027
1 parent a1f8f80
commit 8b51027
Showing 1 changed file with 45 additions and 24 deletions.
diff --git a/src/main/scala/scalismo/io/ActiveShapeModelIO.scala b/src/main/scala/scalismo/io/ActiveShapeModelIO.scala
@@ -34,7 +34,7 @@ object ActiveShapeModelIO {
   object Names {
 
     object Group {
-      val ActiveShapeModel = "/activeShapeModel"
+      val ActiveShapeModel = "activeShapeModel"
       val FeatureExtractor = "featureExtractor"
       val ImagePreprocessor = "imagePreprocessor"
       val Profiles = "profiles"
@@ -55,24 +55,38 @@ object ActiveShapeModelIO {
     }
 
   }
-
   def writeActiveShapeModel(asm: ActiveShapeModel, file: File): Try[Unit] = {
+    for {
+      h5file <- StatisticalModelIOUtils.createFile(file)
+      _ <- writeAcitveShapeModel(asm, h5file, HDFPath("/"))
+      _ <- Try(h5file.close())
+    } yield ()
+  }
 
+  private[io] def writeAcitveShapeModel(asm: ActiveShapeModel, h5file: HDF5Writer, path: HDFPath): Try[Unit] = {
     val pointModel = PointDistributionModel[_3D, TriangleMesh](asm.statisticalModel.gp)
+
+    for {
+      _ <- StatismoIO.writeStatismoPDM(pointModel, h5file, path)
+      _ <- writeAppearanceModel(asm, h5file, path)
+      _ <- h5file.write()
+    } yield ()
+  }
+
+  private def writeAppearanceModel(asm: ActiveShapeModel, h5file: HDF5Writer, path: HDFPath): Try[Unit] = {
+    val asmPath = HDFPath(path, Names.Group.ActiveShapeModel)
+    val fePath = HDFPath(asmPath, Names.Group.FeatureExtractor)
+    val ppPath = HDFPath(asmPath, Names.Group.ImagePreprocessor)
+    val profilesPath = HDFPath(asmPath, Names.Group.Profiles)
+
     for {
-      _ <- StatismoIO.writeStatismoPDM(pointModel, file)
-      h5 <- StatisticalModelIOUtils.openFileForWriting(file)
-      asmPath = HDFPath(Names.Group.ActiveShapeModel)
-      fePath = HDFPath(asmPath, Names.Group.FeatureExtractor)
-      ppPath = HDFPath(asmPath, Names.Group.ImagePreprocessor)
-      profilesPath = HDFPath(asmPath, Names.Group.Profiles)
       // for now, the version information is fixed to 1.0
-      _ <- h5.writeIntAttribute(asmPath, Names.Attribute.MajorVersion, 1)
-      _ <- h5.writeIntAttribute(asmPath, Names.Attribute.MinorVersion, 0)
-      _ <- ImagePreprocessorIOHandlers.save(asm.preprocessor, h5, ppPath)
-      _ <- FeatureExtractorIOHandlers.save(asm.featureExtractor, h5, fePath)
-      _ <- writeProfiles(h5, profilesPath, asm.profiles)
-      _ <- h5.write()
+      _ <- h5file.writeIntAttribute(asmPath, Names.Attribute.MajorVersion, 1)
+      _ <- h5file.writeIntAttribute(asmPath, Names.Attribute.MinorVersion, 0)
+      _ <- ImagePreprocessorIOHandlers.save(asm.preprocessor, h5file, ppPath)
+      _ <- FeatureExtractorIOHandlers.save(asm.featureExtractor, h5file, fePath)
+      _ <- writeProfiles(h5file, profilesPath, asm.profiles)
+      _ <- h5file.write()
     } yield ()
   }
 
@@ -104,10 +118,24 @@ object ActiveShapeModelIO {
     }.flatten
 
   def readActiveShapeModel(fn: File): Try[ActiveShapeModel] = {
+    val modelPath = HDFPath("/")
     for {
-      pdm <- StatismoIO.readStatismoPDM[_3D, TriangleMesh](fn)
       modelReader <- StatisticalModelIOUtils.openFileForReading(fn)
-      asmPath = HDFPath(Names.Group.ActiveShapeModel)
+      pointModel <- StatismoIO.readStatismoPDM[_3D, TriangleMesh](modelReader, modelPath)
+      model <- readActiveShapeModel(modelReader, modelPath, pointModel)
+    } yield model
+  }
+
+  private[io] def readActiveShapeModel(modelReader: StatisticalModelReader,
+                                       modelPath: HDFPath,
+                                       pdm: PointDistributionModel[_3D, TriangleMesh]
+  ): Try[ActiveShapeModel] = {
+    val asmPath = HDFPath(Names.Group.ActiveShapeModel)
+    val fePath = HDFPath(asmPath, Names.Group.FeatureExtractor)
+    val ppPath = HDFPath(asmPath, Names.Group.ImagePreprocessor)
+    val profilesGroup = HDFPath(asmPath, Names.Group.Profiles)
+
+    for {
       asmVersionMajor <- modelReader.readIntAttribute(asmPath, Names.Attribute.MajorVersion)
       asmVersionMinor <- modelReader.readIntAttribute(asmPath, Names.Attribute.MinorVersion)
       _ <- {
@@ -116,17 +144,10 @@ object ActiveShapeModelIO {
           case _ => Failure(new IOException(s"Unsupported ActiveShapeModel version: $asmVersionMajor.$asmVersionMinor"))
         }
       }
-      fePath = HDFPath(asmPath, Names.Group.FeatureExtractor)
-      ppPath = HDFPath(asmPath, Names.Group.ImagePreprocessor)
       preprocessor <- ImagePreprocessorIOHandlers.load(modelReader, ppPath)
-      profilesGroup = HDFPath(asmPath, Names.Group.Profiles)
       featureExtractor <- FeatureExtractorIOHandlers.load(modelReader, fePath)
       profiles <- readProfiles(modelReader, profilesGroup, pdm.reference)
-    } yield {
-      val shapeModel = PointDistributionModel(pdm.gp)
-      ActiveShapeModel(shapeModel, profiles, preprocessor, featureExtractor)
-    }
-
+    } yield ActiveShapeModel(pdm, profiles, preprocessor, featureExtractor)
   }
 
   private[this] def readProfiles(modelReader: StatisticalModelReader,