Merge pull request #24 from hsorby/additional-exports

Additional exports for scaffolds

pyproject.toml

@@ -31,6 +31,7 @@ dependencies = [
   "pennsieve2 >= 0.1.2",
   "cmlibs.zinc >= 4.0.0",
   "cmlibs.utils >= 0.6.0",
+  "cmlibs.exporter >= 0.5.0",
   "scaffoldmaker >= 0.10.0",
   "mbfxml2ex >= 0.7.0",
 ]

src/sparc/client/zinchelper.py

@@ -2,6 +2,8 @@
 import os
 import re
 
+from cmlibs.exporter.vtk import ArgonSceneExporter as VTKExporter
+from cmlibs.exporter.stl import ArgonSceneExporter as STLExporter
 from cmlibs.utils.zinc.field import get_group_list
 from cmlibs.zinc.context import Context
 from cmlibs.zinc.result import RESULT_OK
@@ -103,13 +105,12 @@ def download_files(
         assert response.status_code == 200
         return file_list[0]["name"]
 
-    def get_scaffold_vtk(self, dataset_id, output_file=None):
+    def _get_scaffold(self, dataset_id):
         """
-        Generates a VTK file for the scaffold settings of a dataset.
+        Get a scaffold settings file from the Pennsieve.
 
         Args:
             dataset_id (int): The ID of the dataset to generate the VTK file for.
-            output_file (str): The name of the output VTK file. If not provided, a default name is used.
         """
         scaffold_setting_file = self.download_files(
             limit=1,
@@ -127,10 +128,40 @@ def get_scaffold_vtk(self, dataset_id, output_file=None):
 
         sm = scaffolds.Scaffolds_decodeJSON(c["scaffold_settings"]["scaffoldPackage"])
         sm.generate(self._region)
-        ex = ExportVtk(self._region, "Scaffold VTK export.")
-        if not output_file:
-            output_file = "Scaffold_Creator-settings.vtk"
-        ex.writeFile(output_file)
+
+    def get_scaffold_vtk(self, dataset_id, output_location=None):
+        """
+        Generates a VTK file for the scaffold settings of a dataset.
+
+        Args:
+            dataset_id (int): The ID of the dataset to generate the VTK file for.
+            output_location (str): The output location for the generated VTK file.
+            If not provided, a default of the current working directory is used.
+        """
+        self._get_scaffold(dataset_id)
+
+        if not output_location:
+            output_location = "."
+
+        ex = VTKExporter(output_location, "scaffold")
+        ex.export_vtk_from_scene(self._region.getScene())
+
+    def get_scaffold_stl(self, dataset_id, output_location=None):
+        """
+        Generates an STL file for the scaffold settings of a dataset.
+
+        Args:
+            dataset_id (int): The ID of the dataset to generate the STL file for.
+            output_location (str): The output location for the generated STL file.
+            If not provided, a default of the current working directory is used.
+        """
+        self._get_scaffold(dataset_id)
+
+        if not output_location:
+            output_location = "."
+
+        ex = STLExporter(output_location, "scaffold")
+        ex.export_stl_from_scene(self._region.getScene())
 
     def get_mbf_vtk(self, dataset_id, dataset_file, output_file=None):
         """

tests/test_zinc.py

@@ -10,43 +10,75 @@ def zinc():
     return ZincHelper()
 
 
-def test_get_scaffold_description(zinc):
+def test_export_scaffold_into_vtk_format(zinc):
     # create a temporary output file
-    output_file = os.path.abspath(
-        os.path.join(os.path.dirname(__file__), "resources/scaffold.vtk")
+    output_location = os.path.abspath(
+        os.path.join(os.path.dirname(__file__), "resources/")
     )
 
     # ensure the function returns None if the dataset has no Scaffold_Creator-settings.json file
     invalid_dataset_id = 1000000
     result = None
     with pytest.raises(RuntimeError):
-        result = zinc.get_scaffold_vtk(invalid_dataset_id, output_file)
+        result = zinc.get_scaffold_vtk(invalid_dataset_id, output_location)
     assert result is None
 
     # ensure the function raises an error if the downloaded file is not scaffold_settings file
     dataset_id = 77
     with pytest.raises(AssertionError):
-        zinc.get_scaffold_vtk(dataset_id, output_file)
+        zinc.get_scaffold_vtk(dataset_id, output_location)
 
     # ensure the function generates a VTK file with valid content
     dataset_id = 292
-    zinc.get_scaffold_vtk(dataset_id, output_file)
+    zinc.get_scaffold_vtk(dataset_id, output_location)
+
+    output_file = os.path.join(output_location, "scaffold_root.vtk")
+    assert os.path.exists(output_file)
+    assert os.path.getsize(output_file) > 0
+
+    # Clean up the temporary output file
+    os.remove(output_file)
+
+
+def test_export_scaffold_into_stl_format(zinc):
+    # create a temporary output file
+    output_location = os.path.abspath(
+        os.path.join(os.path.dirname(__file__), "resources/")
+    )
+
+    # ensure the function returns None if the dataset has no Scaffold_Creator-settings.json file
+    invalid_dataset_id = 1000000
+    result = None
+    with pytest.raises(RuntimeError):
+        result = zinc.get_scaffold_stl(invalid_dataset_id, output_location)
+    assert result is None
+
+    # ensure the function raises an error if the downloaded file is not scaffold_settings file
+    dataset_id = 77
+    with pytest.raises(AssertionError):
+        zinc.get_scaffold_stl(dataset_id, output_location)
+
+    # ensure the function generates a VTK file with valid content
+    dataset_id = 292
+    zinc.get_scaffold_stl(dataset_id, output_location)
+
+    output_file = os.path.join(output_location, "scaffold_zinc_graphics.stl")
     assert os.path.exists(output_file)
     assert os.path.getsize(output_file) > 0
 
     # Clean up the temporary output file
     os.remove(output_file)
 
 
-def test_get_scaffold_description_with_default_output_name(zinc):
+def test_export_scaffold_into_vtk_format_with_default_output_location(zinc):
     # ensure the function generates a VTK file with valid content
     dataset_id = 292
     zinc.get_scaffold_vtk(dataset_id)
-    assert os.path.exists("Scaffold_Creator-settings.vtk")
-    assert os.path.getsize("Scaffold_Creator-settings.vtk") > 0
+    assert os.path.exists("scaffold_root.vtk")
+    assert os.path.getsize("scaffold_root.vtk") > 0
 
     # Clean up the temporary output file
-    os.remove("Scaffold_Creator-settings.vtk")
+    os.remove("scaffold_root.vtk")
 
 
 def test_export_mbf_to_vtk(zinc):