e2c17442f525d658434b3da34c52eef9010d7cad
[blender-addons-contrib.git] / io_points_pcd / pcd_utils.py
1 # ##### BEGIN GPL LICENSE BLOCK #####
2 #
3 #  This program is free software; you can redistribute it and/or
4 #  modify it under the terms of the GNU General Public License
5 #  as published by the Free Software Foundation; either version 2
6 #  of the License, or (at your option) any later version.
7 #
8 #  This program is distributed in the hope that it will be useful,
9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11 #  GNU General Public License for more details.
12 #
13 #  You should have received a copy of the GNU General Public License
14 #  along with this program; if not, write to the Free Software Foundation,
15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
16 #
17 # ##### END GPL LICENSE BLOCK #####
18
19 import bpy
20
21 import os
22 import subprocess
23 import tempfile
24
25 from . import pcdparser
26
27
28 def create_and_link_mesh(name, points):
29     """
30     Create a blender mesh and object called name from a list of
31     *points* and link it in the current scene.
32     """
33
34     mesh = bpy.data.meshes.new(name)
35     mesh.from_pydata(points, [], [])
36
37     # update mesh to allow proper display
38     mesh.validate()
39     mesh.update()
40
41     scene = bpy.context.scene
42
43     obj = bpy.data.objects.new(name, mesh)
44     scene.objects.link(obj)
45     obj.select_set(True)
46
47
48 def import_pcd(filepath, name="new_pointcloud"):
49     # check if the file is binary compressed
50     parser = pcdparser.PCDParser.factory(filepath, pcdparser.PointXYZ)
51     parser.onlyParseHeader()
52     is_binary_compressed = parser.datatype == 'BINARY_COMPRESSED'
53
54     # create a temp uncompressed pcd file
55     if (is_binary_compressed):
56         tmpdir = tempfile.TemporaryDirectory()
57         tmpfilepath = tmpdir.name + "blender_pcd_io_tmp.pcd"
58         try:
59             subprocess.call(["pcl_convert_pcd_ascii_binary", filepath, tmpfilepath, "1"])
60         except FileNotFoundError:
61             print("[ERROR] Can't read BINARY COMPRESSED PCD. No pcl_convert_pcd_ascii_binary found! Have you installed libPCL?")
62             return
63         filepath = tmpfilepath
64
65
66     # do actual parsing
67     parser = pcdparser.PCDParser.factory(filepath, pcdparser.PointXYZ)
68     parser.parseFile()
69     points = parser.getPoints()
70
71     blender_points = []
72     for point in points:
73         blender_points.append((point.x, point.y, point.z))
74
75     create_and_link_mesh(name, blender_points)
76
77
78 def export_pcd(filepath):
79     obj = bpy.context.active_object
80
81     # apply object transformation and modifiers
82     mesh = obj.to_mesh(bpy.context.scene, True, "PREVIEW")
83     objmat = obj.matrix_world
84
85     points = []
86     for vert in mesh.vertices:
87         co = objmat * vert.co
88         point = pcdparser.PointXYZ()
89         point.x = co.x
90         point.y = co.y
91         point.z = co.z
92         points.append(point)
93
94     writer = pcdparser.PCDWriter(points)
95     writer.write(filepath)