9f16ca8eb2af948617c37f3bf47f9a90834c8574
[blender-addons-contrib.git] / io_mesh_gwyddion / import_gwyddion.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 import os
21 import re
22 from math import pi, sqrt
23 from mathutils import Vector, Matrix
24 import struct
25
26 # All data for the images. Basically, each variable is a list with a length,
27 # which equals the number of images.
28 # Some of the variables are still not used. However, I keep them for purposes
29 # in future.
30 class AFMData(object):
31     def __init__(self, date, x_size, y_size, x_pixel, y_pixel, x_off, y_off,
32                  voltage, feedback, gain, speed, amplitude, angle, datfile,
33                  channel, unit, z_factor, spec_x_unit, spec_x_label, spec_y_unit,
34                  spec_y_label, spec_y_factor, spec_points, spec_feedback,
35                  spec_acquisition, spec_delay):
36         self.date = date
37         self.x_size = x_size
38         self.y_size = y_size
39         self.x_pixel = x_pixel
40         self.y_pixel = y_pixel
41         self.x_off = x_off
42         self.y_off = y_off
43         self.voltage = voltage
44         self.feedback = feedback
45         self.gain = gain
46         self.speed = speed
47         self.amplitude = amplitude
48         self.angle = angle
49         self.datfile = datfile
50         self.channel = channel
51         self.unit = unit
52         self.z_factor = z_factor
53         self.spec_x_unit = spec_x_unit
54         self.spec_x_label = spec_x_label
55         self.spec_y_unit = spec_y_unit
56         self.spec_y_label = spec_y_label
57         self.spec_y_factor = spec_y_factor
58         self.spec_points = spec_points
59         self.spec_feedback = spec_feedback
60         self.spec_acquisition = spec_acquisition
61         self.spec_delay = spec_delay
62
63
64 # For loading the Gwyddion images. I basically have followed rules described
65 # here: http://gwyddion.net/documentation/user-guide-en/gwyfile-format.html
66 def load_gwyddion_images(data_file, channels):
67
68     if not os.path.isfile(data_file):
69         return False
70
71     AFMdata = AFMData([],[],[],[],[],[],[],
72                       [],[],[],[],[],[],[],
73                       [],[],[],[],[],[],[],
74                       [],[],[],[],[])
75     AFMdata.datfile = data_file
76
77     datafile = open(data_file, 'rb')
78     data = datafile.read()
79     datafile.close()
80
81     # Search the title of each image
82     for a in list(re.finditer(b"data/title\x00", data)):
83
84         pos = a.start()
85         channel_number = int(data[pos-2:pos-1])
86
87         if channels[channel_number] == False:
88             continue
89
90         pos1 = data[pos:].find(b"\x00") + pos + len("\x00") + 1
91         pos2 = data[pos1:].find(b"\x00") + pos1
92
93         channel_name = data[pos1:pos2].decode("utf-8")
94
95         AFMdata.channel.append([channel_number, channel_name])
96
97     # Search important parameters and finally the image data.
98     images = []
99     for a in list(re.finditer(b"/data\x00", data)):
100
101         pos = a.start()
102
103         channel_number = int(data[pos-1:pos])
104
105         if channels[channel_number] == False:
106             continue
107
108         # Find the image size in pixel (x direction)
109         pos1 = data[pos:].find(b"xres") + pos+len("xres")
110         size_x_pixel = struct.unpack("i",data[pos1+2:pos1+4+2])[0]
111
112         # ... the image size in pixel (y direction)
113         pos1 = data[pos:].find(b"yres") + pos+len("yres")
114         size_y_pixel = struct.unpack("i",data[pos1+2:pos1+4+2])[0]
115
116         # ... the real image size (x direction)
117         pos1 = data[pos:].find(b"xreal") + pos+len("xreal")
118         size_x_real = struct.unpack("d",data[pos1+2:pos1+8+2])[0]
119
120         # ... the real image size (y direction)
121         pos1 = data[pos:].find(b"yreal") + pos+len("yreal")
122         size_y_real = struct.unpack("d",data[pos1+2:pos1+8+2])[0]
123
124         # If it is a z image, multiply with 10^9 nm
125         factor = 1.0
126         pos1 = data[pos:].find(b"si_unit_z") + pos
127         unit = data[pos1+34:pos1+36].decode("utf-8")
128         if "m" in unit:
129             factor = 1000000000.0
130
131         # Now, find the image data and store it
132         pos1 = data[pos:].find(b"\x00data\x00") + pos + len("\x00data\x00") + 5
133
134         image = []
135         for i in range(size_y_pixel):
136             line = []
137             for j in range(size_x_pixel):
138                 # The '8' is for the double values
139                 k = pos1 + (i*size_x_pixel+j)   * 8
140                 l = pos1 + (i*size_x_pixel+j+1) * 8
141                 line.append(struct.unpack("d",data[k:l])[0]*factor)
142             image.append(line)
143
144         images.append(image)
145
146         # Note all parameters of the image.
147         AFMdata.x_pixel.append(int(size_x_pixel))
148         AFMdata.y_pixel.append(int(size_y_pixel))
149         AFMdata.x_size.append(size_x_real * 1000000000.0)
150         AFMdata.y_size.append(size_y_real * 1000000000.0)
151
152     return (images, AFMdata)
153
154 # Routine to create the mesh and finally the image
155 def create_mesh(data_list,
156                 AFMdata,
157                 use_smooth,
158                 scale_size,
159                 scale_height,
160                 use_camera,
161                 use_lamp):
162     # This is for the image name.
163     path_list = AFMdata.datfile.strip('/').split('/')
164
165     number_img = len(data_list)
166     image_x_offset_gap = 10.0 * scale_size
167     image_x_all = sum(AFMdata.x_size)*scale_size
168     image_x_offset = -(image_x_all+image_x_offset_gap*(number_img-1)) / 2.0
169
170     # For each image do:
171     for k, data in enumerate(data_list):
172
173         size_x = AFMdata.x_pixel[k]
174         size_y = AFMdata.y_pixel[k]
175
176         image_scale = AFMdata.x_size[k] / float(AFMdata.x_pixel[k])
177         image_scale = image_scale * scale_size
178         image_x_size = AFMdata.x_size[k] * scale_size
179         image_x_offset += image_x_size / 2.0
180
181         image_name = path_list[-1] + "_" + AFMdata.channel[k][1]
182
183         data_mesh = []
184         data_faces = []
185
186         #print("passed - create_mesh ---- 1")
187
188         for i, line in enumerate(data):
189             for j, pixel in enumerate(line):
190
191                # The vertices
192                data_mesh.append(Vector((float(i) * image_scale,
193                                         float(j) * image_scale,
194                                         float(pixel)*scale_height)))
195
196                # The faces
197                if i < size_y-1 and j < size_x-1:
198                    data_faces.append( [size_x*i+j      , size_x*(i+1)+j,
199                                        size_x*(i+1)+j+1, size_x*i+j+1    ])
200
201         #print("passed - create_mesh ---- 2")
202
203         # Build the mesh
204         surface_mesh = bpy.data.meshes.new("Mesh")
205         surface_mesh.from_pydata(data_mesh, [], data_faces)
206         surface_mesh.update()
207         surface = bpy.data.objects.new(image_name, surface_mesh)
208         bpy.context.scene.objects.link(surface)
209         bpy.ops.object.select_all(action='DESELECT')
210         surface.select = True
211
212         bpy.ops.object.origin_set(type='ORIGIN_GEOMETRY')
213         # sum((v.co for v in mesh.vertices), Vector()) / len(mesh.vertices)
214
215         if use_smooth:
216             for polygon in surface.data.polygons:
217                 polygon.use_smooth = True
218
219         surface.location = Vector((0.0, image_x_offset, 0.0))
220         image_x_offset += image_x_size / 2.0 + image_x_offset_gap
221
222         #print("passed - create_mesh ---- 3")
223
224
225
226     object_center_vec = Vector((0.0,0.0,0.0))
227     object_size = (sum(AFMdata.x_size) * scale_size
228                    +image_x_offset_gap * (len(data_list)-1))
229
230     # ------------------------------------------------------------------------
231     # CAMERA AND LAMP
232     camera_factor = 20.0
233
234     # If chosen a camera is put into the scene.
235     if use_camera == True:
236
237         # Assume that the object is put into the global origin. Then, the
238         # camera is moved in x and z direction, not in y. The object has its
239         # size at distance sqrt(object_size) from the origin. So, move the
240         # camera by this distance times a factor of camera_factor in x and z.
241         # Then add x, y and z of the origin of the object.
242         object_camera_vec = Vector((sqrt(object_size) * camera_factor,
243                                     0.0,
244                                     sqrt(object_size) * camera_factor))
245         camera_xyz_vec = object_center_vec + object_camera_vec
246
247         # Create the camera
248         current_layers=bpy.context.scene.layers
249         camera_data = bpy.data.cameras.new("A_camera")
250         camera_data.lens = 45
251         camera_data.clip_end = 50000.0
252         camera = bpy.data.objects.new("A_camera", camera_data)
253         camera.location = camera_xyz_vec
254         camera.layers = current_layers
255         bpy.context.scene.objects.link(camera)
256
257         # Here the camera is rotated such it looks towards the center of
258         # the object. The [0.0, 0.0, 1.0] vector along the z axis
259         z_axis_vec             = Vector((0.0, 0.0, 1.0))
260         # The angle between the last two vectors
261         angle                  = object_camera_vec.angle(z_axis_vec, 0)
262         # The cross-product of z_axis_vec and object_camera_vec
263         axis_vec               = z_axis_vec.cross(object_camera_vec)
264         # Rotate 'axis_vec' by 'angle' and convert this to euler parameters.
265         # 4 is the size of the matrix.
266         camera.rotation_euler  = Matrix.Rotation(angle, 4, axis_vec).to_euler()
267
268         # Rotate the camera around its axis by 90° such that we have a nice
269         # camera position and view onto the object.
270         bpy.ops.object.select_all(action='DESELECT')
271         camera.select_set(True)
272         bpy.ops.transform.rotate(value=(90.0*2*pi/360.0),
273                                  axis=object_camera_vec,
274                                  constraint_axis=(False, False, False),
275                                  constraint_orientation='GLOBAL',
276                                  mirror=False, proportional='DISABLED',
277                                  proportional_edit_falloff='SMOOTH',
278                                  proportional_size=1, snap=False,
279                                  snap_target='CLOSEST', snap_point=(0, 0, 0),
280                                  snap_align=False, snap_normal=(0, 0, 0),
281                                  release_confirm=False)
282
283     # Here a lamp is put into the scene, if chosen.
284     if use_lamp == True:
285
286         # This is the distance from the object measured in terms of %
287         # of the camera distance. It is set onto 50% (1/2) distance.
288         lamp_dl = sqrt(object_size) * 15 * 0.5
289         # This is a factor to which extend the lamp shall go to the right
290         # (from the camera  point of view).
291         lamp_dy_right = lamp_dl * (3.0/4.0)
292
293         # Create x, y and z for the lamp.
294         object_lamp_vec = Vector((lamp_dl,lamp_dy_right,lamp_dl))
295         lamp_xyz_vec = object_center_vec + object_lamp_vec
296
297         # Create the lamp
298         current_layers=bpy.context.scene.layers
299         lamp_data = bpy.data.lamps.new(name="A_lamp", type="POINT")
300         lamp_data.distance = 5000.0
301         lamp_data.energy = 3.0
302         lamp_data.shadow_method = 'RAY_SHADOW'
303         lamp = bpy.data.objects.new("A_lamp", lamp_data)
304         lamp.location = lamp_xyz_vec
305         lamp.layers = current_layers
306         bpy.context.scene.objects.link(lamp)
307
308         bpy.context.scene.world.light_settings.use_ambient_occlusion = True
309         bpy.context.scene.world.light_settings.ao_factor = 0.1