Merge pull request #59 from joaori/master

Adding support for writing BlobNode (imagery) with test case for e57 file cloning

src/pye57/libe57_wrapper.cpp

@@ -197,6 +197,9 @@ PYBIND11_MODULE(libe57, m) {
     cls_StructureNode.def("set", [](StructureNode &node, const std::string &pathName, StringNode &n){
         node.set(pathName, n);
     }, "pathName"_a, "n"_a);
+    cls_StructureNode.def("set", [](StructureNode &node, const std::string &pathName, BlobNode &n){
+        node.set(pathName, n);
+    }, "pathName"_a, "n"_a);
     cls_StructureNode.def(py::init<const e57::Node &>(), "n"_a);
     cls_StructureNode.def("isRoot", &StructureNode::isRoot);
     cls_StructureNode.def("parent", &StructureNode::parent);

src/pye57/utils.py

@@ -1,5 +1,3 @@
-from typing import Type
-
 from pye57 import libe57
 from pye57.libe57 import NodeType
 
@@ -39,4 +37,81 @@ def convert_spherical_to_cartesian(rae):
         range_cos_phi * np.cos(theta),
         range_cos_phi * np.sin(theta),
         range_ * np.sin(phi)
-    ), axis=1)
+    ), axis=1)
+
+
+def copy_node(node, dest_image):
+    compressed_node_pairs = []
+    blob_node_pairs = []
+
+    out_node = None
+    # 'Element' Types
+    if (isinstance(node, libe57.FloatNode)):
+        out_node = libe57.FloatNode(
+            dest_image,
+            value=node.value(),
+            precision=node.precision(),
+            minimum=node.minimum(),
+            maximum=node.maximum())
+
+    elif (isinstance(node, libe57.IntegerNode)):
+        out_node = libe57.IntegerNode(
+            dest_image,
+            value=node.value(),
+            minimum=node.minimum(),
+            maximum=node.maximum())
+
+    elif (isinstance(node, libe57.ScaledIntegerNode)):
+        out_node = libe57.ScaledIntegerNode(
+            dest_image,
+            node.rawValue(),
+            minimum=node.minimum(),
+            maximum=node.maximum(),
+            scale=node.scale(),
+            offset=node.offset())
+
+    elif (isinstance(node, libe57.StringNode)):
+        out_node = libe57.StringNode(
+            dest_image,
+            node.value())
+
+    elif (isinstance(node, libe57.BlobNode)):
+        out_node = libe57.BlobNode(dest_image, node.byteCount())
+        blob_node_pairs.append({ 'in': node, 'out': out_node })
+
+    # 'Container' Types
+    elif (isinstance(node, libe57.CompressedVectorNode)):
+        in_prototype = libe57.StructureNode(node.prototype())
+        out_prototype, _, _  = copy_node(in_prototype, dest_image)
+        out_codecs, _, _ = copy_node(node.codecs(), dest_image)
+
+        out_node = libe57.CompressedVectorNode(dest_image, out_prototype, out_codecs)
+
+        compressed_node_pairs.append({
+            'in': node,
+            'out': out_node
+        })
+
+    elif isinstance(node, libe57.StructureNode):
+        out_node = libe57.StructureNode(dest_image)
+        for i in range(node.childCount()):
+            in_child = get_node(node, i)
+            in_child_name = in_child.elementName()
+            out_child, out_child_compressed_node_pairs, out_child_blob_node_pairs = copy_node(in_child, dest_image)
+
+            out_node.set(in_child_name, out_child)
+            compressed_node_pairs.extend(out_child_compressed_node_pairs)
+            blob_node_pairs.extend(out_child_blob_node_pairs)
+
+    elif isinstance(node, libe57.VectorNode):
+        out_node = libe57.VectorNode(dest_image, allowHeteroChildren=node.allowHeteroChildren())
+        for i in range(node.childCount()):
+            in_child = get_node(node, i)
+            in_child_name = f'{i}'
+            out_child, out_child_compressed_node_pairs, out_child_blob_node_pairs = copy_node(in_child, dest_image)
+
+            out_node.append(out_child)
+            compressed_node_pairs.extend(out_child_compressed_node_pairs)
+            blob_node_pairs.extend(out_child_blob_node_pairs)
+
+    return out_node, compressed_node_pairs, blob_node_pairs

tests/test_main.py

@@ -6,6 +6,7 @@
 
 import pye57
 from pye57 import libe57
+from pye57.utils import get_node, copy_node
 
 try:
     from exceptions import WindowsError
@@ -37,10 +38,18 @@ def delete_retry(path):
 def e57_path():
     return sample_data("test.e57")
 
+
 @pytest.fixture
 def e57_spherical_path():
     return sample_data("testSpherical.e57")
 
+
+@pytest.fixture
+def e57_with_data_and_images_path():
+    # From http://www.libe57.org/data.html
+    return sample_data("pumpAVisualReferenceImage.e57")
+
+
 @pytest.fixture
 def temp_e57_write(request):
     path = sample_data("test_write.e57")
@@ -302,3 +311,111 @@ def test_read_color_absent(e57_path):
 def test_scan_position(e57_path):
     e57 = pye57.E57(e57_path)
     assert np.allclose(e57.scan_position(3), np.array([[3.01323456e+05, 5.04260184e+06, 1.56040279e+01]]))
+
+
+BUFFER_TYPES = {
+    libe57.FloatNode: 'd',
+    libe57.IntegerNode: 'l',
+    libe57.ScaledIntegerNode: 'd'
+}
+
+
+def make_buffer(node: libe57.Node, capacity: int):
+    node_type = type(node)
+    if node_type not in BUFFER_TYPES:
+        raise ValueError("Unsupported field type!")
+
+    buffer_type = BUFFER_TYPES[node_type]
+
+    np_array = np.empty(capacity, buffer_type)
+    buffer = libe57.SourceDestBuffer(
+        node.destImageFile(), node.elementName(), np_array, capacity, True, True)
+    return np_array, buffer
+
+
+def make_buffers(node: libe57.StructureNode, capacity: int):
+    data = {}
+    buffers = libe57.VectorSourceDestBuffer()
+    for i in range(node.childCount()):
+        field = get_node(node, i)
+        d, b = make_buffer(field, capacity)
+        data[field.elementName()] = d
+        buffers.append(b)
+    return data, buffers
+
+
+def copy_compressed_vector_data(in_node: libe57.Node, out_node: libe57.Node):
+    chunk_size = 100000
+
+    in_prototype = libe57.StructureNode(in_node.prototype())
+    out_prototype = libe57.StructureNode(out_node.prototype())
+
+    in_data, in_buffers = make_buffers(in_prototype, chunk_size)
+    out_data, out_buffers = make_buffers(out_prototype, chunk_size)
+
+    in_reader = in_node.reader(in_buffers)
+    out_writer = out_node.writer(out_buffers)
+
+    n_points = in_node.childCount()
+    current_index = 0
+    while current_index != n_points:
+        current_chunk = min(n_points - current_index, chunk_size)
+
+        in_reader.read()
+        for field in in_data:
+            out_data[field][:current_chunk] = in_data[field][:current_chunk]
+
+        out_writer.write(current_chunk)
+
+        current_index += current_chunk
+
+    in_reader.close()
+    out_writer.close()
+
+
+def copy_blob_data(in_node, out_node):
+    chunk_size = 100000
+
+    byte_count = in_node.byteCount()
+    blob_buffer = np.empty(chunk_size, np.ubyte)
+    current_index = 0
+    while current_index != byte_count:
+        current_chunk = min(byte_count - current_index, chunk_size)
+
+        in_node.read(blob_buffer, current_index, current_chunk)
+        out_node.write(blob_buffer, current_index, current_chunk)
+
+        current_index += current_chunk
+
+
+def test_clone_e57(e57_with_data_and_images_path, temp_e57_write):
+
+    in_image = libe57.ImageFile(e57_with_data_and_images_path, "r")
+    out_image = libe57.ImageFile(temp_e57_write, "w")
+
+    for i in range(in_image.extensionsCount()):
+        out_image.extensionsAdd(
+            in_image.extensionsPrefix(i),
+            in_image.extensionsUri(i)
+        )
+
+    in_root = in_image.root()
+    out_root = out_image.root()
+
+    compressed_node_pairs = []
+    for i in range(in_root.childCount()):
+        in_child = get_node(in_root, i)
+        in_child_name = in_child.elementName()
+        out_child, out_child_compressed_node_pairs, out_child_blob_node_pairs = copy_node(in_child, out_image)
+
+        out_root.set(in_child_name, out_child)
+        compressed_node_pairs.extend(out_child_compressed_node_pairs)
+
+    for compressed_node_pair in compressed_node_pairs:
+        copy_compressed_vector_data(compressed_node_pair['in'], compressed_node_pair['out'])
+
+    for blob_node_pair in out_child_blob_node_pairs:
+        copy_blob_data(blob_node_pair['in'], blob_node_pair['out'])
+
+    in_image.close()
+    out_image.close()