👨🏿‍🔬 Refactor polygon and STL handling

Author: ZenanH

src/FastPointQuery.jl

@@ -33,43 +33,8 @@ function __init__()
     end
 end
 
-struct STLInfo2D
-    mesh        ::Py
-    vmin        ::Vector
-    vmax        ::Vector
-    py_vertices ::Py
-    py_triangles::Py
-    stl_path    ::String
-end
-
-struct STLInfo3D
-    mesh        ::Py
-    vmin        ::Vector
-    vmax        ::Vector
-    py_vertices ::Py
-    py_triangles::Py
-    stl_path    ::String
-end
-
-struct QueryPolygon
-    polygon::Py
-    coord  ::AbstractVector
-end
-
-function QueryPolygon(pypoly::Py)
-    @pyexec """
-    def py_tmp(poly, np):
-        rings = []
-        exterior = np.array(poly.exterior.coords)
-        rings.append(exterior)
-        for interior in poly.interiors:
-            hole = np.array(interior.coords)
-            rings.append(hole)
-        return rings
-    """ => py_tmp
-    coord = py2ju(Vector, py_tmp(pypoly, np))
-    return QueryPolygon(pypoly, coord)
-end
+include(joinpath(@__DIR__, "fileio/geojson.jl"))
+include(joinpath(@__DIR__, "fileio/stl.jl"))
 
 include(joinpath(@__DIR__, "utils.jl"))
 include(joinpath(@__DIR__, "polygon.jl"))

src/discretization/_polygon.jl

@@ -1,3 +1,9 @@
+#==========================================================================================+
+| TABLE OF CONTENTS:                                                                       |
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
+|  - get_pts                                                                               |
++==========================================================================================#
+
 export get_pts
 
 """
@@ -66,15 +72,9 @@ pts = get_pts(stl_data, 0.1; fill=true) # returns points in a filled grid patter
 """
 function get_pts(stl_data::STLInfo2D, h::Real; fill::Bool=true)
     h > 0 || error("h must be positive")
-    vertices = stl_data.py_vertices
-    triangles = stl_data.py_triangles
-    @pyexec """
-    def py_tmp(vertices, triangles, shapely):
-        tris = vertices[triangles]
-        tris_2d = tris[:, :, :2]
-        polygons = [shapely.Polygon(pts) for pts in tris_2d]
-        return shapely.unary_union(polygons)
-    """ => py_tmp
-    polygon = QueryPolygon(py_tmp(vertices, triangles, shapely))
+    triangle_coords_2d = stl_data.py_vertices[stl_data.py_triangles]
+    tris2d = shapely.polygons(triangle_coords_2d)
+    region = shapely.unary_union(tris2d)
+    polygon = QueryPolygon(region)
     return get_pts(polygon, h; fill=fill)
 end

src/fileio/geojson.jl

@@ -0,0 +1,102 @@
+#==========================================================================================+
+| TABLE OF CONTENTS:                                                                       |
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
+| struct:                                                                                  |
+|  - QueryPolygon                                                                          |
++------------------------------------------------------------------------------------------+
+| function:                                                                                |
+|  - read_polygon                                                                          |
+|  - write_polygon                                                                         |
++==========================================================================================#
+
+export QueryPolygon
+export read_polygon, write_polygon
+
+struct QueryPolygon
+    polygon::Py
+    coord  ::AbstractVector
+end
+
+function QueryPolygon(pypoly::Py)
+    @pyexec """
+    def py_tmp(poly, np):
+        rings = []
+        exterior = np.array(poly.exterior.coords)
+        rings.append(exterior)
+        for interior in poly.interiors:
+            hole = np.array(interior.coords)
+            rings.append(hole)
+        return rings
+    """ => py_tmp
+    coord = py2ju(Vector, py_tmp(pypoly, np))
+    return QueryPolygon(pypoly, coord)
+end
+
+"""
+    write_polygon(polygon::QueryPolygon, file_path::String)
+
+Description:
+---
+Write the polygon to a GeoJSON file. The `polygon` is an instance of `QueryPolygon`, and
+`file_path` is the path where the GeoJSON file will be saved.
+
+Example:
+---
+```julia
+polygon_xy = [0 0; 1 0; 1 1; 0 1]' # 2xN array
+poly = get_polygon(polygon_xy, ratio=1) # poly.polygon to visualize
+write_polygon(poly, "/path/to/polygon.geojson")
+# or
+write_polygon(poly, "/path/to/polygon") # will automatically add .geojson
+"""
+function write_polygon(polygon::QueryPolygon, file_path::String)
+    if !endswith(lowercase(file_path), ".geojson")
+        file_path *= ".geojson"
+    end
+    @pyexec """
+    def py_tmp(poly, filename, json, shapely):
+        feature = {
+            "type": "Feature",
+            "geometry": shapely.geometry.mapping(poly),
+            "properties": {}
+        }
+        geojson = {
+            "type": "FeatureCollection",
+            "features": [feature]
+        }
+        with open(filename, "w") as f:
+            json.dump(geojson, f, indent=2)
+    """ => py_tmp
+    py_tmp(polygon.polygon, file_path, pyjson, shapely)
+    @info """geojson file saved at:
+    $file_path
+    """
+end
+
+"""
+    read_polygon(file_path::String)
+
+Description:
+---
+Read a polygon from a GeoJSON file. The `file_path` is the path to the GeoJSON file.
+
+Example:
+---
+```julia
+polygon = read_polygon("/path/to/polygon.geojson")
+```
+"""
+function read_polygon(file_path::String)
+    isfile(file_path) || error("file_path should be a valid file path")
+    endswith(lowercase(file_path), ".geojson") || error("file_path is not a .geojson file")
+    @pyexec """
+    def py_tmp(filename, json, shapely):
+        with open(filename, "r") as f:
+            data = json.load(f)
+        feature = data["features"][0]
+        geometry = feature["geometry"]
+        polygon = shapely.geometry.shape(geometry)
+        return polygon
+    """ => py_tmp
+    return QueryPolygon(py_tmp(file_path, pyjson, shapely))
+end

src/fileio/stl.jl

@@ -0,0 +1,101 @@
+#==========================================================================================+
+| TABLE OF CONTENTS:                                                                       |
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
+| struct:                                                                                  |  
+| - STLInfo2D                                                                              |
+| - STLInfo3D                                                                              |
++------------------------------------------------------------------------------------------+
+| function:                                                                                |
+|  - readSTL2D                                                                             |
+|  - readSTL3D                                                                             |
++==========================================================================================#
+
+export STLInfo2D, STLInfo3D
+export readSTL2D, readSTL3D
+
+struct STLInfo2D
+    mesh        ::Py
+    vmin        ::Vector
+    vmax        ::Vector
+    py_vertices ::Py
+    py_triangles::Py
+    stl_path    ::String
+end
+
+function Base.show(io::IO, info::STLInfo2D)
+    println(io, "STLInfo2D:")
+    println(io, "  Path     : ", info.stl_path)
+    println(io, "  Vertices : ", info.py_vertices.shape[0])
+    println(io, "  Triangles: ", info.py_triangles.shape[0])
+    println(io, "  X-Y Min  : (", info.vmin[1], ", ", info.vmin[2], ")")
+    println(io, "  X-Y Max  : (", info.vmax[1], ", ", info.vmax[2], ")")
+end
+
+struct STLInfo3D
+    mesh        ::Py
+    vmin        ::Vector
+    vmax        ::Vector
+    py_vertices ::Py
+    py_triangles::Py
+    stl_path    ::String
+end
+
+function Base.show(io::IO, info::STLInfo3D)
+    println(io, "STLInfo3D:")
+    println(io, "  Path     : ", info.stl_path)
+    println(io, "  Vertices : ", info.py_vertices.shape[0])
+    println(io, "  Triangles: ", info.py_triangles.shape[0])
+    println(io, "  X-Y-Z Min: (", info.vmin[1], ", ", info.vmin[2], ", ", info.vmin[3], ")")
+    println(io, "  X-Y-Z Max: (", info.vmax[1], ", ", info.vmax[2], ", ", info.vmax[3], ")")
+end
+
+"""
+    readSTL2D(stl_file::String)
+
+Description:
+---
+Read a STL file as a 2D object (z=0) containing the mesh, bounding box, vertices, and 
+triangles.
+
+Example:
+---
+```julia
+stl_info = readSTL2D("path/to/your/file.stl")
+```
+"""
+function readSTL2D(stl_file)
+    isfile(stl_file) || error("stl_file should be a valid file path")
+    mesh = o3d.io.read_triangle_mesh(stl_file)
+    rows = pyslice(nothing, nothing, nothing)
+    cols = pyslice(nothing, 2, nothing)
+    vertices = np.asarray(mesh.vertices)[rows, cols]
+    triangles = np.asarray(mesh.triangles)
+    aabb = mesh.get_axis_aligned_bounding_box()
+    vmin = py2ju(Vector, aabb.get_min_bound())[1:2]
+    vmax = py2ju(Vector, aabb.get_max_bound())[1:2]
+    return STLInfo2D(mesh, vmin, vmax, vertices, triangles, stl_file) 
+end
+
+"""
+    readSTL3D(stl_file::String)
+
+Description:
+---
+Read a STL file as a 3D object containing the mesh, bounding box, vertices, and triangles.
+
+Example:
+---
+```julia
+stl_info = readSTL3D("path/to/your/file.stl")
+```
+"""
+function readSTL3D(stl_file)
+    isfile(stl_file) || error("stl_file should be a valid file path")
+    mesh = o3d.io.read_triangle_mesh(stl_file)
+    vertices = np.asarray(mesh.vertices)
+    triangles = np.asarray(mesh.triangles)
+    aabb = mesh.get_axis_aligned_bounding_box()
+    vmin = py2ju(Vector, aabb.get_min_bound())
+    vmax = py2ju(Vector, aabb.get_max_bound())
+    return STLInfo3D(mesh, vmin, vmax, vertices, triangles, stl_file) 
+end

src/polygon.jl

@@ -33,7 +33,7 @@ function pip_query(
     edge   ::Bool=false
 )
     func = edge ? shapely.intersects_xy : shapely.contains_xy
-    return pyconvert(Bool, func(polygon.py_xy, px, py).item())
+    return py2ju(Bool, func(polygon.py_xy, px, py).item())
 end
 
 """
@@ -65,13 +65,14 @@ function pip_query(
     py_points = shapely.points(points')
 
     # check if the particle is inside the polygon
+    println("\e[1;36m[start]:\e[0m querying points inside the polygon")
     if edge
         mask = shapely.covers(polygon.polygon, py_points)
     else
         mask = shapely.within(py_points, polygon.polygon)
     end
 
-    return pyconvert(Vector{Bool}, mask)
+    return PyArray(mask)
 end
 
 """
@@ -99,25 +100,17 @@ function pip_query(
 )
     size(points, 1) == 2 || error("points must be a 2xN array")
     py_points = shapely.points(points')
+    triangle_coords_2d = stl_model.py_vertices[stl_model.py_triangles]
+    tris2d = shapely.polygons(triangle_coords_2d)
+    region = shapely.unary_union(tris2d)
+
+    # check if the particle is inside the polygon
+    println("\e[1;36m[start]:\e[0m querying points inside the STL model")
     if edge
-        @pyexec """
-        def py_pip(vertices, triangles, py_points, shapely):
-            triangle_coords_2d = vertices[triangles][:, :, :2]
-            tris2d = shapely.polygons(triangle_coords_2d)
-            region = shapely.unary_union(tris2d)
-            mask = shapely.covers(region, py_points)
-            return mask
-        """ => py_pip
+        mask = shapely.covers(region, py_points)
     else
-        @pyexec """
-        def py_pip(vertices, triangles, py_points, shapely):
-            triangle_coords_2d = vertices[triangles][:, :, :2]
-            tris2d = shapely.polygons(triangle_coords_2d)
-            region = shapely.unary_union(tris2d)
-            mask = shapely.within(py_points, region)
-            return mask
-        """ => py_pip
+        mask = shapely.within(py_points, region)
     end
-    tmp = py_pip(stl_model.py_vertices, stl_model.py_triangles, py_points, shapely)
-    return pyconvert(Vector{Bool}, tmp)
+
+    return PyArray(mask)
 end