Add initial batch processing and do some tidying.

src/stack3d/formats/unet.py

@@ -0,0 +1,136 @@
+""" Format conversion for Russ Bate's unet format.
+"""
+
+from __future__ import print_function
+import inspect
+import os
+import csv
+from subprocess import call
+from argparse import ArgumentParser
+import scipy.io
+import networkx as nx
+import stack3d.formats
+
+
+def load_skeleton(path):
+
+    """
+    Load the skeleton from a pickle
+    """
+
+    # Delayed import so script can be run with both Python 2 and 3
+    from unet_core.vessel_analysis import VesselTree
+
+    v = VesselTree()
+    v.load_skeleton(path)
+    return v.skeleton
+
+
+def skeleton_to_vtp(path):
+
+    temp_output_path = os.path.dirname(os.path.realpath(path)) + "/temp_network_csv.csv"
+    script_path = inspect.getfile(stack3d.formats.unet).replace(".pyc", ".py")
+    launch_args = " --input '" + path + "' --output '" + temp_output_path + "' --format csv --field diameter"
+    call("python3 " + script_path + launch_args, shell=True)
+
+    # Avoid vtk import when method not used to reduce dependency overhead
+    import vtk
+    polydata = vtk.vtkPolyData()
+    points = vtk.vtkPoints()
+    cells = vtk.vtkCellArray()
+    data = vtk.vtkDoubleArray()
+
+    with open(temp_output_path, 'rb') as csvfile:
+        reader = csv.reader(csvfile, delimiter=',', quotechar='|')
+        counter = 0
+        first = 0
+        for row in reader:
+            if first == 0:
+                first += 1
+                continue
+            point0 = [float(x) for x in row[0:3]]
+            point1 = [float(x) for x in row[3:6]]
+            diameter = float(row[6])
+            points.InsertNextPoint(point0)
+            points.InsertNextPoint(point1)
+            line = vtk.vtkLine()
+            line.GetPointIds().SetId(0, counter)
+            line.GetPointIds().SetId(1, counter+1)
+            counter += 2
+            cells.InsertNextCell(line)
+            data.InsertNextTuple1(diameter)
+
+        polydata.SetPoints(points)
+        polydata.SetLines(cells)
+        polydata.GetCellData().SetScalars(data)
+
+        clean = vtk.vtkCleanPolyData()
+        clean.SetInputData(polydata)
+        clean.Update()
+
+        return clean.GetOutput()
+
+
+def skeleton_to_matlab(skeleton, field, output_path):
+
+    """
+    Convert the skeleton's network description to matlab sparse matrix format.
+    Each mat file corresponds to a connected component
+    """
+
+#     if os.path.isabs(output_path) and not os.path.exists(output_path):
+#         os.makedirs(output_path)
+
+    for idx, component in enumerate(skeleton.components):
+        for n, nbrs in component.graph.adjacency_iter():
+            for nbr, edict in nbrs.items():
+                branch = edict[0]["branch"]
+                component.graph.add_edge(n, nbr, weight=branch.diameter)
+        scipy_mat = nx.to_scipy_sparse_matrix(component.graph)
+        output = {"skeleton": scipy_mat}
+        output_file = output_path + "/skeleton_" + field + "_Comp_" + str(idx) + ".mat"
+        scipy.io.savemat(output_file, output)
+
+def skeleton_to_csv(skeleton, output_path):
+
+    """
+    Convert the skeleton's network description to csv format
+    """
+
+#     if os.path.isabs(output_path) and not os.path.exists(output_path):
+#         os.makedirs(output_path)
+
+    f = open(output_path, "w")
+    f.write("P0 - x, P0 - y, P0 - z, P1 - x, P1 - y, P1 - z, diameter")
+
+    for _, _, branch in skeleton.skeleton_branch_iter():
+        for idx, p1 in enumerate(branch.points):
+            if idx > 0:
+                p0 = branch.points[idx-1]
+                p0_s = str(p0.x) + "," + str(p0.y) + "," + str(p0.z)
+                p1_s = str(p1.x) + "," + str(p1.y) + "," + str(p1.z)
+                f.write(p0_s + "," + p1_s + "," + str(branch.diameter) + "\n")
+    f.close()
+
+
+if __name__ == "__main__":
+
+    parser = ArgumentParser()
+    parser.add_argument('--input', type=str, help='Input skeleton file.')
+    parser.add_argument('--output', type=str, help='Output directory.')
+    parser.add_argument('--format', type=str,
+                        help='Output format.',
+                        choices=['csv', 'mat'],
+                        default='mat')
+    parser.add_argument('--field', type=str,
+                        help='Output field.',
+                        choices=['diameter', "length", "tortuosity"],
+                        default='diameter')
+    args = parser.parse_args()
+
+    skeleton = load_skeleton(args.input)
+    if "csv" in args.format:
+        skeleton_to_csv(skeleton, args.output)
+    elif "mat" in args.format:
+        skeleton_to_matlab(skeleton, args.field, args.output)
+