SimpleCrownGenerator

  1# coding: utf-8
  2# Author: Jianbo Qi
  3# Date: 2020.4.16
  4# Generate tree crowns with randomly distributed leaves
  5import numpy as np
  6from scipy.linalg import expm, norm
  7import math, sys
  8import trimesh
  9import random
 10from scipy.spatial import cKDTree
 11import alphashape
 12from AlphaShapeCGAL import AlphaShapeCGAL
 13import tempfile
 14import shutil
 15
 16
 17def agd_amp_max(sig_l, sig_r, gamma):
 18    tmp = math.sqrt(math.gamma(1 / gamma) / math.gamma(3 / gamma))
 19    beta_l = sig_l * tmp
 20    beta_r = sig_r * tmp
 21    A = gamma / ((beta_l + beta_r) * math.gamma(1 / gamma))
 22    return A
 23
 24
 25def agd(x, sig_l, sig_r, gamma=2.0, m=0.0):
 26    """
 27    Asymmetric gaussian distribution
 28    """
 29    tmp = math.sqrt(math.gamma(1 / gamma) / math.gamma(3 / gamma))
 30    beta_l = sig_l * tmp
 31    beta_r = sig_r * tmp
 32    A = gamma / ((beta_l + beta_r) * math.gamma(1 / gamma))
 33    if x <= m:
 34        y = A * math.exp(-(-(x - m) / beta_l) ** gamma)
 35    else:
 36        y = A * math.exp(-((x - m) / beta_r) ** gamma)
 37    return y
 38
 39
 40class LAD:
 41    UNIFORM = "Uniform"
 42    SPHERICAL = "Spherical"
 43    ERECTOPHILE = "Erectophile"
 44    EXTREMOPHILE = "Extremophile"
 45    PLANOPHILE = "Planophile"
 46    PLAGIOPHILE = "Plagiophile"
 47
 48
 49class CrownShape:
 50    CUBE = "Cube"
 51    ELLIPSOID = "Ellipsoid"
 52    CYLINDER = "Cylinder"
 53    CONE = "Cone"
 54    ASYMMETRIC_GAUSSIAN = "Asymmeric_Gaussian"
 55    # USER_DEFINED_CURVE = "USER_DEFINED_CURVE"
 56
 57
 58class LeafShape:
 59    SQUARE = "Square"
 60    DISK = "Disk"
 61
 62
 63class CrownGenerator:
 64    def __init__(self):
 65        self.leaf_angle_dist = LAD.SPHERICAL
 66        self.leaf_volume_density = 1.0
 67        self.crown_diameter_SN = 1
 68        self.crown_diameter_EW = 1
 69        self.crown_height = 3
 70        self.crown_shape = CrownShape.ELLIPSOID
 71        self.leaf_shape = LeafShape.SQUARE
 72        self.leaf_num_triangles = 12
 73        self.single_leaf_area = 0.01
 74
 75        # generate crown boundary only
 76        self.generate_crown_only = False
 77        self.clumped_volume = 0  # If the leaf is turbid, and clumping is less than 1, store the clumped volume
 78        self.total_wood_area = 0
 79
 80        # trunk
 81        self.has_trunk = False
 82        self.trunk_height = 3
 83        self.dbh = 0.2
 84        # This value divides trunk into two segments, for applications such as forest fire severity
 85        self.trunk_division_height = 0
 86
 87        # If generate branches, when trunk is generated
 88        self.has_branches = False
 89        self.branch_vertical_angle = 45  # degree, angle between [0, 1, 0]
 90        self.vertical_distance_between_branch = 0.8  # vertical distance between branches, starting from trunk height
 91        self.num_branches_each_height = 3  # The number of branches at each branch height
 92        self.branch_radius_scale = 0.33  # the scale factor for branch base radius, compared with parent trunk
 93        self.random_angle_range = 20  # the random angle range to offset vertical and horizontal angles
 94
 95        # clumping
 96        # [0, 1] 1 for random, 0 for total clumping
 97        self.clumping_factor = 1.0
 98        # [0, 1) clustering leaf to branch top factor: 1 around branch top, 0 around whole branch
 99        self.leaf_distance_factor_to_branch_root = 0.3
100
101        # trunk for asymmetric gaussian
102        self.sig_upper = 2
103        self.sig_lower = 4
104        self.gamma_shape = 6
105
106        # curve for user defined, [(Height, Crown_Radius)]
107        # self.user_curve = [(0, 0.4), (0.2, 0.5)]
108
109        self._num_of_leaves = 100
110        self._num_sigma = 2.0
111
112    def get_crown_volume(self):
113        volume = 0.0
114        if self.crown_shape == CrownShape.CUBE:
115            volume = self.crown_diameter_SN * self.crown_diameter_EW * self.crown_height
116        elif self.crown_shape == CrownShape.ELLIPSOID:
117            volume = 4 * math.pi * (0.5 * self.crown_diameter_EW) * (0.5 * self.crown_diameter_SN) * (
118                        0.5 * self.crown_height) / 3.0
119        elif self.crown_shape == CrownShape.CYLINDER:
120            volume = math.pi * (0.5 * self.crown_diameter_EW) * (0.5 * self.crown_diameter_SN) * self.crown_height
121        elif self.crown_shape == CrownShape.CONE:
122            volume = math.pi * (0.5 * self.crown_diameter_EW) * (0.5 * self.crown_diameter_SN) * self.crown_height / 3.0
123        elif self.crown_shape == CrownShape.ASYMMETRIC_GAUSSIAN:
124            num_sigma = self._num_sigma
125            n_bins = 499
126            x = np.linspace(-num_sigma * self.sig_lower, num_sigma * self.sig_upper, n_bins + 1)
127            step = self.crown_height / n_bins
128            amp_max = agd_amp_max(self.sig_lower, self.sig_upper, self.gamma_shape)
129            y = []
130            for xx in x:
131                yy = agd(xx, self.sig_lower, self.sig_upper, gamma=self.gamma_shape)
132                y.append(self.crown_diameter_SN * 0.5 * yy / amp_max)
133            for i in range(len(x) - 1):
134                s = (1.0 / 3.0) * math.pi * step * (y[i] ** 2 + y[i + 1] ** 2 + y[i] * y[i + 1])
135                volume += s
136        return volume
137
138    def __tree_var_to_asymmetric_gaussian_xy(self, h, dist_trunk):
139        """
140        This function compute the x value, given the height value in the crown
141        """
142        num_sigma = self._num_sigma
143        min_x, max_x = -num_sigma * self.sig_lower, num_sigma * self.sig_upper
144        x = min_x + (max_x - min_x) * (h / self.crown_height)
145        amp_max = agd_amp_max(self.sig_lower, self.sig_upper, self.gamma_shape)
146        y = dist_trunk / (0.5 * self.crown_diameter_SN) * amp_max
147        return x, y, amp_max
148
149    # considering trunk, if leaf position is within radius, it should be regenerated.
150    def consider_by_trunk_radius(self, x, y, z, leaf_length):
151        if not self.has_trunk:
152            return [x, y, z]  # if no trunk, directly return the original xyz
153        else:
154            h = self.trunk_height + y
155            trunk_radius = self.get_trunk_radius_at_height(h)
156            dist_to_center = np.sqrt(x * x + z * z)
157            while dist_to_center < trunk_radius + leaf_length:
158                if self.crown_shape == CrownShape.CUBE:
159                    x = -0.5 * self.crown_diameter_EW + np.random.rand() * self.crown_diameter_EW
160                    z = -0.5 * self.crown_diameter_SN + np.random.rand() * self.crown_diameter_SN
161                    y = np.random.rand() * self.crown_height
162                elif self.crown_shape == CrownShape.ELLIPSOID:
163                    phi = np.random.rand() * np.pi * 2  # randomly azimuth angle
164                    theta = np.arccos(2 * np.random.rand() - 1)  # randomly zenith angle
165                    r = np.cbrt(np.random.rand())
166                    x = r * self.crown_diameter_EW * 0.5 * np.sin(theta) * np.cos(phi)
167                    z = r * self.crown_diameter_SN * 0.5 * np.sin(theta) * np.sin(phi)
168                    y = r * self.crown_height * 0.5 * np.cos(theta) + self.crown_height * 0.5
169                elif self.crown_shape == CrownShape.CYLINDER:
170                    phi = np.random.rand() * np.pi * 2  # randomly azimuth angle
171                    r = np.sqrt(np.random.rand())
172                    x = r * self.crown_diameter_EW * 0.5 * np.cos(phi)
173                    z = r * self.crown_diameter_SN * 0.5 * np.sin(phi)
174                    y = np.random.rand() * self.crown_height
175                elif self.crown_shape == CrownShape.CONE:
176                    phi = np.random.rand() * np.pi * 2  # randomly azimuth angle
177                    h = self.crown_height * (np.random.rand()) ** (1 / 3)
178                    rnd = np.random.rand()
179                    rew = (0.5 * self.crown_diameter_EW / self.crown_height) * h * np.sqrt(rnd)
180                    rsn = (0.5 * self.crown_diameter_SN / self.crown_height) * h * np.sqrt(rnd)
181                    x = rew * np.cos(phi)
182                    z = rsn * np.sin(phi)
183                    y = self.crown_height - h
184                elif self.crown_shape == CrownShape.ASYMMETRIC_GAUSSIAN:
185                    while True:
186                        y = self.crown_height * np.random.rand()
187                        x = -0.5 * self.crown_diameter_EW + np.random.rand() * self.crown_diameter_EW
188                        z = -0.5 * self.crown_diameter_SN + np.random.rand() * self.crown_diameter_SN
189                        dist_to_center1 = np.sqrt(x * x + z * z)
190                        asym_x, asym_y, amp_max = self.__tree_var_to_asymmetric_gaussian_xy(y, dist_to_center1)
191                        pred_y = agd(asym_x, self.sig_lower, self.sig_upper, self.gamma_shape)
192                        if asym_y <= pred_y:
193                            break
194                dist_to_center = np.sqrt(x * x + z * z)
195            return [x, y, z]
196
197    def _generate_leaf_pos(self):
198        leaf_length = self._calculate_leaf_length()
199        all_pos = []
200        if self.crown_shape == CrownShape.CUBE:
201            for i in range(self._num_of_leaves):
202                x = -0.5 * self.crown_diameter_EW + np.random.rand() * self.crown_diameter_EW
203                z = -0.5 * self.crown_diameter_SN + np.random.rand() * self.crown_diameter_SN
204                y = np.random.rand() * self.crown_height
205                all_pos.append(self.consider_by_trunk_radius(x, y, z, leaf_length))
206
207        elif self.crown_shape == CrownShape.ELLIPSOID:
208            for i in range(self._num_of_leaves):
209                phi = np.random.rand() * np.pi * 2  # randomly azimuth angle
210                theta = np.arccos(2 * np.random.rand() - 1)  # randomly zenith angle
211                r = np.cbrt(np.random.rand())
212                x = r * self.crown_diameter_EW * 0.5 * np.sin(theta) * np.cos(phi)
213                z = r * self.crown_diameter_SN * 0.5 * np.sin(theta) * np.sin(phi)
214                y = r * self.crown_height * 0.5 * np.cos(theta) + self.crown_height * 0.5
215                all_pos.append(self.consider_by_trunk_radius(x, y, z, leaf_length))
216        elif self.crown_shape == CrownShape.CYLINDER:
217            for i in range(self._num_of_leaves):
218                phi = np.random.rand() * np.pi * 2  # randomly azimuth angle
219                r = np.sqrt(np.random.rand())
220                x = r * self.crown_diameter_EW * 0.5 * np.cos(phi)
221                z = r * self.crown_diameter_SN * 0.5 * np.sin(phi)
222                y = np.random.rand() * self.crown_height
223                all_pos.append(self.consider_by_trunk_radius(x, y, z, leaf_length))
224        elif self.crown_shape == CrownShape.CONE:
225            for i in range(self._num_of_leaves):
226                phi = np.random.rand() * np.pi * 2  # randomly azimuth angle
227                h = self.crown_height * (np.random.rand()) ** (1 / 3)
228                rnd = np.random.rand()
229                rew = (0.5 * self.crown_diameter_EW / self.crown_height) * h * np.sqrt(rnd)
230                rsn = (0.5 * self.crown_diameter_SN / self.crown_height) * h * np.sqrt(rnd)
231                x = rew * np.cos(phi)
232                z = rsn * np.sin(phi)
233                y = self.crown_height - h
234                all_pos.append(self.consider_by_trunk_radius(x, y, z, leaf_length))
235        elif self.crown_shape == CrownShape.ASYMMETRIC_GAUSSIAN:
236            while True:
237                y = self.crown_height * np.random.rand()
238                x = -0.5 * self.crown_diameter_EW + np.random.rand() * self.crown_diameter_EW
239                z = -0.5 * self.crown_diameter_SN + np.random.rand() * self.crown_diameter_SN
240                dist_to_center = np.sqrt(x * x + z * z)
241                asym_x, asym_y, amp_max = self.__tree_var_to_asymmetric_gaussian_xy(y, dist_to_center)
242                pred_y = agd(asym_x, self.sig_lower, self.sig_upper, self.gamma_shape)
243                if asym_y <= pred_y:
244                    all_pos.append(self.consider_by_trunk_radius(x, y, z, leaf_length))
245                if len(all_pos) == self._num_of_leaves:
246                    break
247        # elif self.crown_shape == CrownShape.USER_DEFINED_CURVE:
248        #     while True:
249        #         y = self.crown_height * np.random.rand()
250        #         x = -0.5 * self.crown_diameter_EW + np.random.rand() * self.crown_diameter_EW
251        #         z = -0.5 * self.crown_diameter_SN + np.random.rand() * self.crown_diameter_SN
252        #         dist_to_center = np.sqrt(x * x + z * z)
253        return all_pos
254
255    def _generate_leaf_normal(self):
256        '''
257        Sphrical: 球形状
258        Uniform:统一型
259        Planophile：平面型
260        Erectophile：竖直型
261        Plagiophile：倾斜型
262        Extremophile：极端型
263        '''
264        all_normals = []
265        if self.leaf_angle_dist == LAD.SPHERICAL:
266            for i in range(self._num_of_leaves):
267                phi = np.random.rand() * np.pi * 2  # randomly azimuth angle
268                theta = np.arccos(np.random.rand())  # randomly zenith angle
269                rx = -np.sin(theta) * np.cos(phi)
270                rz = np.sin(theta) * np.sin(phi)
271                ry = np.cos(theta)
272                all_normals.append([rx, ry, rz])
273            # f = open(r"D:\LESS\simulations\LESSMedium\SimpleHorizontalLayerRAMI\HOM03_ERE_scene.def")
274            # for line in f:
275            #     arr = list(map(lambda x: float(x), line.split()))
276            #     all_normals.append([-arr[4], arr[6], arr[5]])
277        if self.leaf_angle_dist == LAD.UNIFORM:
278            for i in range(self._num_of_leaves):
279                phi = np.random.rand() * np.pi * 2  # randomly azimuth angle
280                theta = np.random.rand() * np.pi * 0.5  # randomly zenith angle
281                rx = -np.sin(theta) * np.cos(phi)
282                rz = np.sin(theta) * np.sin(phi)
283                ry = np.cos(theta)
284                all_normals.append([rx, ry, rz])
285        if self.leaf_angle_dist == LAD.PLANOPHILE:
286            while True:
287                theta = np.random.rand() * np.pi * 0.5
288                y = (2 + 2 * np.cos(2 * theta)) / np.pi
289                rnd_y = np.random.rand() * 4 / np.pi
290                if rnd_y <= y:  # accept
291                    phi = np.random.rand() * np.pi * 2  # randomly azimuth angle
292                    rx = -np.sin(theta) * np.cos(phi)
293                    rz = np.sin(theta) * np.sin(phi)
294                    ry = np.cos(theta)
295                    all_normals.append([rx, ry, rz])
296                    if len(all_normals) == self._num_of_leaves:
297                        break
298        if self.leaf_angle_dist == LAD.ERECTOPHILE:  # 竖直型
299            while True:
300                theta = np.random.rand() * np.pi * 0.5
301                y = (2 - 2 * np.cos(2 * theta)) / np.pi
302                rnd_y = np.random.rand() * 4 / np.pi
303                if rnd_y <= y:  # accept
304                    phi = np.random.rand() * np.pi * 2  # randomly azimuth angle
305                    rx = -np.sin(theta) * np.cos(phi)
306                    rz = np.sin(theta) * np.sin(phi)
307                    ry = np.cos(theta)
308                    all_normals.append([rx, ry, rz])
309                    if len(all_normals) == self._num_of_leaves:
310                        break
311        if self.leaf_angle_dist == LAD.PLAGIOPHILE:  # 倾斜型
312            while True:
313                theta = np.random.rand() * np.pi * 0.5
314                y = (2 - 2 * np.cos(4 * theta)) / np.pi
315                rnd_y = np.random.rand() * 4 / np.pi
316                if rnd_y <= y:  # accept
317                    phi = np.random.rand() * np.pi * 2  # randomly azimuth angle
318                    rx = -np.sin(theta) * np.cos(phi)
319                    rz = np.sin(theta) * np.sin(phi)
320                    ry = np.cos(theta)
321                    all_normals.append([rx, ry, rz])
322                    if len(all_normals) == self._num_of_leaves:
323                        break
324        if self.leaf_angle_dist == LAD.EXTREMOPHILE:  # 极端
325            while True:
326                theta = np.random.rand() * np.pi * 0.5
327                y = (2 + 2 * np.cos(4 * theta)) / np.pi
328                rnd_y = np.random.rand() * 4 / np.pi
329                if rnd_y <= y:  # accept
330                    phi = np.random.rand() * np.pi * 2  # randomly azimuth angle
331                    rx = -np.sin(theta) * np.cos(phi)
332                    rz = np.sin(theta) * np.sin(phi)
333                    ry = np.cos(theta)
334                    all_normals.append([rx, ry, rz])
335                    if len(all_normals) == self._num_of_leaves:
336                        break
337        return all_normals
338
339    @staticmethod
340    def _M(axis, theta):
341        return expm(np.cross(np.eye(3), axis / norm(axis) * theta))
342
343    def _generate_single_leaf(self, leaf_pos, leaf_normal, leaf_length):
344        leaf_side_length = leaf_length
345        # generate leaf according to normal
346        if self.leaf_shape == LeafShape.SQUARE:
347            # First, find a arbitrary vector which is not parallel with normal
348            # and do cross product, the resulting vector is in leaf plane
349            arv = np.array([0, 0, 1])
350            v = np.cross(leaf_normal, arv)
351            v = v / np.linalg.norm(v)
352            p1 = leaf_side_length * np.sqrt(2) * 0.5 * v + leaf_pos
353            v = np.cross(leaf_normal, v)
354            p2 = leaf_side_length * np.sqrt(2) * 0.5 * v + leaf_pos
355            v = np.cross(leaf_normal, v)
356            p3 = leaf_side_length * np.sqrt(2) * 0.5 * v + leaf_pos
357            v = np.cross(leaf_normal, v)
358            p4 = leaf_side_length * np.sqrt(2) * 0.5 * v + leaf_pos
359            return [p1, p2, p3, p4]
360        elif self.leaf_shape == LeafShape.DISK:  # leaf_side_length is leaf radius
361            # First, find a arbitrary vector which is not parallel with normal
362            # and do cross product, the resulting vector is in leaf plane
363            pts = []
364            rot_rad = np.pi * 2 / self.leaf_num_triangles
365            arv = np.array([0, 0, 1])
366            v0 = np.cross(leaf_normal, arv)
367            v0 = v0 / np.linalg.norm(v0)
368            p1 = leaf_side_length * v0 + leaf_pos  # 第一个点
369            pts.append(p1)
370            for i in range(0, self.leaf_num_triangles - 1):
371                m0 = CrownGenerator._M(np.array(leaf_normal), rot_rad * (i + 1))
372                newv = np.dot(m0, v0)
373                newv = newv / np.linalg.norm(newv)
374                newp = leaf_side_length * newv + leaf_pos  # 第一个点
375                pts.append(newp)
376            return pts
377        else:
378            print("Leaf shape is not supported.")
379
380    def _calculate_leaf_length(self):
381        if self.leaf_shape == LeafShape.SQUARE:
382            return np.sqrt(self.single_leaf_area)
383        if self.leaf_shape == LeafShape.DISK:
384            theta = np.pi * 2 / self.leaf_num_triangles
385            tri_leaf = self.single_leaf_area / self.leaf_num_triangles
386            leaf_radius = np.sqrt(2 * tri_leaf / np.sin(theta))
387            return leaf_radius
388
389    def __compute_leaf_numbers(self):
390        crown_volume = self.get_crown_volume()
391        tot_area = crown_volume*self.leaf_volume_density
392        self._num_of_leaves = int(tot_area / self.single_leaf_area)
393        print("Crown volume:", "%.4f" % crown_volume)
394        print("Number of leaves:", self._num_of_leaves)
395        sys.stdout.flush()
396
397    def __generate_crown_boundary(self, h_of_first_branch):
398        mesh = None
399        bound_box = np.array([self.crown_diameter_EW, self.crown_height, self.crown_diameter_SN])
400        center_pt = np.array([0, h_of_first_branch+0.5*self.crown_height, 0])
401        if self.crown_shape == CrownShape.CUBE:
402            mesh = trimesh.creation.box(extents=bound_box)
403            mesh.apply_translation(center_pt)
404        elif self.crown_shape == CrownShape.ELLIPSOID:
405            mesh = trimesh.creation.icosphere()
406            mesh.apply_scale(0.5 * bound_box)
407            mesh.apply_translation(center_pt)
408        elif self.crown_shape == CrownShape.CYLINDER:
409            mesh = trimesh.creation.cylinder(1, 1)
410            mesh.vertices[:, [1, 2]] = mesh.vertices[:, [2, 1]]
411            mesh.faces[:, [1, 2]] = mesh.faces[:, [2, 1]]
412            mesh.apply_scale([0.5 * bound_box[0], bound_box[1], 0.5 * bound_box[2]])
413            mesh.apply_translation(center_pt)
414        elif self.crown_shape == CrownShape.CONE:
415            mesh = trimesh.creation.cone(1, 1)
416            mesh.vertices[:, [1, 2]] = mesh.vertices[:, [2, 1]]
417            mesh.faces[:, [1, 2]] = mesh.faces[:, [2, 1]]
418            mesh.apply_scale([0.5 * bound_box[0], bound_box[1], 0.5 * bound_box[2]])
419            center_pt = np.array([0, h_of_first_branch, 0])
420            mesh.apply_translation(center_pt)
421        elif self.crown_shape == CrownShape.ASYMMETRIC_GAUSSIAN:
422            num_sigma = self._num_sigma
423            n_bins = 199
424            min_bound = -num_sigma * self.sig_lower
425            max_bound = num_sigma * self.sig_upper
426            x_locals = np.linspace(min_bound, max_bound, n_bins + 1)
427            step = self.crown_height / n_bins
428            amp_max = agd_amp_max(self.sig_lower, self.sig_upper, self.gamma_shape)
429            xy = []
430            for x_local in x_locals:
431                y_local = agd(x_local, self.sig_lower, self.sig_upper, gamma=self.gamma_shape)
432                y_global = self.crown_diameter_SN * 0.5 * y_local / amp_max
433                x_global = h_of_first_branch + self.crown_height*(x_local - min_bound) / (max_bound - min_bound)
434                xy.append([y_global, x_global])
435            xy[0][0]=0
436            xy[-1][0]=0
437            mesh = trimesh.creation.revolve(np.array(xy))
438            mesh.vertices[:, [1, 2]] = mesh.vertices[:, [2, 1]]
439            mesh.faces[:, [1, 2]] = mesh.faces[:, [2, 1]]
440        return mesh
441
442    def get_trunk_radius_at_height(self, h):
443        dbh_radis = 0.5 * self.dbh
444        ground_radius = dbh_radis * (self.trunk_height + self.crown_height) / (
445                    self.trunk_height + self.crown_height - 1.3)
446        if self.trunk_height + self.crown_height <= 1.3:
447            ground_radius = dbh_radis
448        division_radius = ground_radius * (self.trunk_height + self.crown_height - h) / (
449                    self.trunk_height + self.crown_height)
450        if division_radius < 0:
451            division_radius = 0
452        return division_radius
453
454    def get_trunk_surface_area(self):
455        dbh_radis = 0.5 * self.dbh
456        ground_radius = dbh_radis * (self.trunk_height + self.crown_height) / (
457                self.trunk_height + self.crown_height - 1.3)
458        if self.trunk_height + self.crown_height <= 1.3:
459            ground_radius = dbh_radis
460        h = self.trunk_height + self.crown_height
461        return math.pi*ground_radius*math.sqrt(h*h+ground_radius*ground_radius)
462
463    def get_trunk_volume(self):
464        dbh_radis = 0.5 * self.dbh
465        ground_radius = dbh_radis * (self.trunk_height + self.crown_height) / (
466                self.trunk_height + self.crown_height - 1.3)
467        if self.trunk_height + self.crown_height <= 1.3:
468            ground_radius = dbh_radis
469        return math.pi*ground_radius*ground_radius*(self.trunk_height + self.crown_height)/3
470
471
472    def generate_crown(self, dist_obj_path):
473        f_out = open(dist_obj_path, "w")
474        trunk_start_index = 0
475        h_of_first_branch = 0
476        if self.has_trunk:
477            h_of_first_branch = self.trunk_height
478        crown_mesh = None
479        # trunk
480        branch_points = []
481        branch_ids = []
482        if self.has_trunk:
483            points_lower = []
484            points_upper = []
485            points_ground = []
486            points_division = []
487            f_out.write("g trunk\n")
488            dbh_radis = 0.5 * self.dbh
489            ground_radius = dbh_radis*(self.trunk_height + self.crown_height)/(self.trunk_height + self.crown_height-1.3)
490            if self.trunk_height + self.crown_height <= 1.3:
491                print("Total tree height is less than 1.3, using DBH as tree base radius.")
492                ground_radius = dbh_radis
493            division_radius = ground_radius*(self.trunk_height + self.crown_height-self.trunk_division_height)/(self.trunk_height + self.crown_height)
494            if self.trunk_height + self.crown_height-self.trunk_division_height <=0:
495                self.trunk_division_height = self.trunk_height + self.crown_height
496                division_radius = 0.001
497
498            for i in range(0, 360, 15):
499                angle_rad = np.radians(i)
500                x_ground = -ground_radius * np.cos(angle_rad)
501                z_ground = ground_radius * np.sin(angle_rad)
502                x_division = -division_radius * np.cos(angle_rad)
503                z_division = division_radius * np.sin(angle_rad)
504                x_top = -0.001 * np.cos(angle_rad)
505                z_top = 0.001 * np.sin(angle_rad)
506                points_lower.append([x_division, self.trunk_division_height, z_division])
507                points_upper.append([x_top, self.trunk_height + self.crown_height, z_top])
508                if self.trunk_division_height > 0:
509                    points_ground.append([x_ground, 0, z_ground])
510                    points_division.append([x_division, self.trunk_division_height, z_division])
511            points_lower.append(points_lower[0])
512            points_upper.append(points_upper[0])
513            if self.trunk_division_height > 0:
514                points_ground.append(points_ground[0])
515                points_division.append(points_division[0])
516            for i in range(0, 24):  # 12个面元
517                f_out.write("v %.4f %.4f %.4f\n" % (points_lower[i][0], points_lower[i][1], points_lower[i][2]))
518                f_out.write(
519                    "v %.4f %.4f %.4f\n" % (points_lower[i + 1][0], points_lower[i + 1][1], points_lower[i + 1][2]))
520                f_out.write(
521                    "v %.4f %.4f %.4f\n" % (points_upper[i + 1][0], points_upper[i + 1][1], points_upper[i + 1][2]))
522                f_out.write("v %.4f %.4f %.4f\n" % (points_upper[i][0], points_upper[i][1], points_upper[i][2]))
523                f_out.write(
524                    "f %d %d %d %d\n" % (trunk_start_index + 1,
525                                         trunk_start_index + 2,
526                                         trunk_start_index + 3,
527                                         trunk_start_index + 4))
528                trunk_start_index += 4
529            #  trunk top
530            for i in range(0, 24):  # 12 top triangles
531                f_out.write(
532                    "v %.4f %.4f %.4f\n" % (points_upper[i][0], points_upper[i][1], points_upper[i][2]))
533                f_out.write(
534                    "v %.4f %.4f %.4f\n" % (points_upper[i + 1][0], points_upper[i + 1][1], points_upper[i + 1][2]))
535                f_out.write(
536                    "v %.4f %.4f %.4f\n" % (0, self.trunk_height, 0))
537                f_out.write(
538                    "f %d %d %d\n" % (trunk_start_index + 1,
539                                      trunk_start_index + 2,
540                                      trunk_start_index + 3))
541                trunk_start_index += 3
542
543            if self.trunk_division_height > 0:
544                f_out.write("\ng trunk_bottom\n")
545                for i in range(0, 24):  # 12个面元
546                    f_out.write("v %.4f %.4f %.4f\n" % (points_ground[i][0], points_ground[i][1], points_ground[i][2]))
547                    f_out.write(
548                        "v %.4f %.4f %.4f\n" % (points_ground[i + 1][0], points_ground[i + 1][1], points_ground[i + 1][2]))
549                    f_out.write(
550                        "v %.4f %.4f %.4f\n" % (points_division[i + 1][0], points_division[i + 1][1], points_division[i + 1][2]))
551                    f_out.write("v %.4f %.4f %.4f\n" % (points_division[i][0], points_division[i][1], points_division[i][2]))
552                    f_out.write(
553                        "f %d %d %d %d\n" % (trunk_start_index + 1,
554                                             trunk_start_index + 2,
555                                             trunk_start_index + 3,
556                                             trunk_start_index + 4))
557                    trunk_start_index += 4
558
559            branch_surface_area = 0
560            branch_volume = 0
561            if self.has_branches:
562                f_out.write("\ng branch\n")
563                if crown_mesh is None:
564                    crown_mesh = self.__generate_crown_boundary(h_of_first_branch)
565                branch_vertical_angle_center = self.branch_vertical_angle/180*math.pi
566                branch_horizontal_angle_interval = 2*math.pi/self.num_branches_each_height
567                num_branch_heights = int(self.crown_height/self.vertical_distance_between_branch)
568                index = 0
569                branch_index = 0
570                for i in range(num_branch_heights):
571                    branch_height = h_of_first_branch + i*self.vertical_distance_between_branch + 0.0001
572                    h_radius = self.get_trunk_radius_at_height(branch_height)
573                    angle_start = 0
574                    if index % 2 == 1:
575                        angle_start = 0.5 * branch_horizontal_angle_interval
576                    for j in range(self.num_branches_each_height):  # for each branch
577                        random_angle_range = self.random_angle_range/180*math.pi
578                        branch_vertical_angle = branch_vertical_angle_center + random_angle_range*(random.random()*2-1)
579                        horizontal_angle = angle_start + j*branch_horizontal_angle_interval + random_angle_range*(random.random()*2-1)
580                        x = -np.sin(branch_vertical_angle)*np.cos(horizontal_angle)
581                        z = np.sin(branch_vertical_angle)*np.sin(horizontal_angle)
582                        y = np.cos(branch_vertical_angle)
583                        ray_origins = np.array([[0, branch_height, 0]])
584                        ray_directions = np.array([[x, y, z]])
585                        locations, index_ray, index_tri = crown_mesh.ray.intersects_location(
586                            ray_origins=ray_origins,
587                            ray_directions=ray_directions)
588                        dist = np.linalg.norm(locations[0] - ray_origins[0])
589                        if self.clumping_factor < 1:
590                            point_seg = int(dist/0.01)
591                            for k in range(int(point_seg*self.leaf_distance_factor_to_branch_root), point_seg, 1):
592                                # offset to remove h_of_first_branch to match with leaf all_pos
593                                branch_points.append(ray_origins[0] + (k/point_seg)*dist*ray_directions[0]-np.array([0, h_of_first_branch, 0]))
594                                if self.generate_crown_only:
595                                    branch_ids.append(branch_index)
596                        mesh = trimesh.creation.cone(1, 1)
597                        mesh.vertices[:, [1, 2]] = mesh.vertices[:, [2, 1]]
598                        mesh.faces[:, [1, 2]] = mesh.faces[:, [2, 1]]
599                        mesh.apply_scale(
600                            [self.branch_radius_scale * h_radius, dist, self.branch_radius_scale * h_radius])
601                        branch_surface_area += (mesh.area-math.pi*math.pow(self.branch_radius_scale * h_radius, 2))
602                        branch_volume += mesh.volume
603                        mesh.faces += trunk_start_index
604                        trunk_start_index += len(mesh.vertices)
605                        R1 = trimesh.transformations.rotation_matrix(branch_vertical_angle, [0, 0, 1])
606                        R2 = trimesh.transformations.rotation_matrix(horizontal_angle, [0, 1, 0])
607                        R = trimesh.transformations.concatenate_matrices(R2, R1)
608                        mesh.apply_transform(R)
609                        center_pt = np.array([0, branch_height, 0])
610                        mesh.apply_translation(center_pt)
611                        obj = trimesh.exchange.obj.export_obj(mesh)
612                        f_out.write("\n"+obj)
613                        branch_index += 1
614                    index += 1
615            trunk_surface = self.get_trunk_surface_area()
616            trunk_volume = self.get_trunk_volume()
617            self.total_wood_area = trunk_surface + branch_surface_area
618            print("Total wood surface area: %.4f" % (trunk_surface + branch_surface_area))
619            print(" - Total branch surface area: %.4f" % branch_surface_area)
620            print(" - Total trunk surface area: %.4f" % trunk_surface)
621            print("Total wood volume: %.4f" % (branch_volume + trunk_volume))
622            print(" - Total trunk volume: %.4f" % trunk_volume)
623            print(" - Total branch volume: %.4f" % branch_volume)
624
625        if not self.generate_crown_only:
626            self.__compute_leaf_numbers()
627            all_pos = self._generate_leaf_pos()
628            if len(all_pos) > 0:
629                if self.clumping_factor < 1 and self.has_branches and self.has_trunk:
630                    kdtree = cKDTree(branch_points)
631                    near_dists, p_idxs = kdtree.query(all_pos, k=1)
632                    for i in range(len(p_idxs)):
633                        p_idx = p_idxs[i]
634                        branch_point = branch_points[p_idx]
635                        leaf_point = np.array(all_pos[i])
636                        # move the leaf towards the branch
637                        vec_dir = branch_point - leaf_point
638                        dist = np.linalg.norm(vec_dir)
639                        vec_dir = vec_dir/dist
640                        new_leaf_point = leaf_point + vec_dir*dist*(1-self.clumping_factor)
641                        all_pos[i] = new_leaf_point
642
643                all_normal = self._generate_leaf_normal()
644                f_out.write("\ng leaves\n")
645                leaf_length = self._calculate_leaf_length()
646                for i in range(self._num_of_leaves):
647                    leaf_pos = all_pos[i]
648                    leaf_normal = all_normal[i]
649                    points = self._generate_single_leaf(leaf_pos, leaf_normal, leaf_length)
650                    if i % 10000 == 0:
651                        f_out.flush()
652                    if self.leaf_shape == LeafShape.SQUARE:
653                        p1, p2, p3, p4 = points
654                        f_out.write("v %.4f %.4f %.4f\n" % (p1[0], p1[1] + h_of_first_branch, p1[2]))
655                        f_out.write("v %.4f %.4f %.4f\n" % (p2[0], p2[1] + h_of_first_branch, p2[2]))
656                        f_out.write("v %.4f %.4f %.4f\n" % (p3[0], p3[1] + h_of_first_branch, p3[2]))
657                        f_out.write("v %.4f %.4f %.4f\n" % (p4[0], p4[1] + h_of_first_branch, p4[2]))
658                        f_out.write(
659                            "f %d %d %d %d\n" % (trunk_start_index+1, trunk_start_index + 2, trunk_start_index + 3, trunk_start_index + 4))
660                        trunk_start_index += 4
661                    if self.leaf_shape == LeafShape.DISK:
662                        points.append(leaf_pos)
663                        tot_pt_each_leaf = len(points)
664                        for j in range(0, len(points)):
665                            f_out.write("v %.4f %.4f %.4f\n" % (points[j][0], points[j][1] + h_of_first_branch, points[j][2]))
666                        for j in range(0, self.leaf_num_triangles - 1):
667                            f_out.write(
668                                "f %d %d %d\n" % (trunk_start_index+ j + 1,
669                                                  trunk_start_index+ j + 2,
670                                                  trunk_start_index+ tot_pt_each_leaf))
671                        f_out.write(
672                            "f %d %d %d\n" % (trunk_start_index+self.leaf_num_triangles,
673                                              trunk_start_index + 1,
674                                              trunk_start_index + tot_pt_each_leaf))
675                        trunk_start_index += tot_pt_each_leaf
676        else:  # generate boundary
677            crown_volume = self.get_crown_volume()
678            print("Corwn volume:", "%.4f" % crown_volume)
679            if crown_mesh is None:
680                crown_mesh = self.__generate_crown_boundary(h_of_first_branch)
681            f_out.write("\ng leaf_boundary\n")
682            if self.clumping_factor < 1 and self.has_branches and self.has_trunk:
683                branch_points=branch_points+np.array([0, h_of_first_branch, 0])
684                kdtree = cKDTree(branch_points)
685                num_samples = int(crown_volume * 300)
686                sample_points = trimesh.sample.volume_mesh(crown_mesh, num_samples)
687                near_dists, p_idxs = kdtree.query(sample_points, k=1)
688                sample_points_bran_idxs=np.zeros(sample_points.shape[0], dtype=int)
689                for i in range(len(p_idxs)):
690                    p_idx = p_idxs[i]
691                    branch_point = branch_points[p_idx]
692                    sample_points_bran_idxs[i] = branch_ids[p_idx]
693                    leaf_point = np.array(sample_points[i])
694                    # move the leaf towards the branch
695                    vec_dir = branch_point - leaf_point
696                    dist = np.linalg.norm(vec_dir)
697                    vec_dir = vec_dir / dist
698                    new_leaf_point = leaf_point + vec_dir * dist * (1 - self.clumping_factor)
699                    sample_points[i] = new_leaf_point
700                tot_vol = 0
701                tmp_dir = tempfile.mkdtemp()
702                alpha_cg = AlphaShapeCGAL(tmp_dir, True)
703                for idx in np.unique(sample_points_bran_idxs):
704                    single_samole_points = sample_points[sample_points_bran_idxs == idx]
705                    # alpha_shape = alphashape.alphashape(single_samole_points, 0.1)
706                    if single_samole_points.shape[0] > 3:
707                        alpha_shape = alpha_cg.alphashape(single_samole_points, 5)
708                        if isinstance(alpha_shape, trimesh.base.Trimesh):
709                            tot_vol += alpha_shape.volume
710                            alpha_shape.faces += trunk_start_index
711                            obj = trimesh.exchange.obj.export_obj(alpha_shape)
712                            f_out.write("\n"+obj)
713                            trunk_start_index += len(alpha_shape.vertices)
714                alpha_cg.free_library()
715                shutil.rmtree(tmp_dir)
716                print("Crown volume after clumping:", "%.4f" % tot_vol)
717                self.clumped_volume = tot_vol
718
719            else:
720                crown_mesh.faces += trunk_start_index
721                obj = trimesh.exchange.obj.export_obj(crown_mesh)
722                f_out.write(obj)
723                trunk_start_index += len(crown_mesh.vertices)
724                self.clumped_volume = crown_volume
725            f_out.close()
726
727
728if __name__ == "__main__":
729    import argparse
730
731    parser = argparse.ArgumentParser()
732    parser.add_argument("--crown_shape", help="Option: Ellipsoid, Cube, Cylinder, Cone, Asymmeric_Gaussian", type=str, default="Ellipsoid")
733    parser.add_argument('--crown_height', help="Unit: m", type=float, default=4)
734    parser.add_argument("--crown_diameter_sn", help="Unit: m", type=float, default=2)
735    parser.add_argument("--crown_diameter_ew", help="Unit: m", type=float, default=2)
736    parser.add_argument("--trunk_height", help="Unit: m", type=float, default=3)
737    parser.add_argument("--dbh", help="Diameter at breast height, Unit: m", type=float, default="0.2")
738    parser.add_argument("--lad", help="Option:Spherical, Uniform, Planophile, Plagiophile, Erectophile, "
739                                      "Extremophile", type=str, default="Spherical")
740    parser.add_argument("--lower_sigma", help="Lower shape of the crown for Asymmeric_Gaussian", type=float, default=0.1)
741    parser.add_argument("--upper_sigma", help="Upper shape of the crown for Asymmeric_Gaussian", type=float,
742                        default=1)
743    parser.add_argument("--gamma_shape", help="Lower shape of the crown for Asymmeric_Gaussian", type=float,
744                        default=1)
745    parser.add_argument("--lvd", help="", type=float, default=1.0)
746    parser.add_argument("--leaf_shape", help="Option: Square, Disk", type=str, default="Square")
747    parser.add_argument("--polygon_sides", help="Only for disk leaf", type=int, default=6)
748    parser.add_argument("--single_leaf_area", help="Unit: Square meter", type=float, default=0.0025)
749    parser.add_argument("--generate_crown_only", help="Generate crown boundary only", action="store_true")
750
751    parser.add_argument("--has_branches", help="Has branches", action="store_true")
752    parser.add_argument("--branch_vertical_angle", help="Angle between branch and trunk", type=float, default=45)
753    parser.add_argument("--vertical_distance_between_branch", help="vertical_distance_between_branch", type=float, default=0.8)
754    parser.add_argument("--num_branches_each_height", help="num_branches_each_height", type=int,
755                        default=3)
756    parser.add_argument("--branch_radius_scale", help="branch_radius_scale", type=float,
757                        default=0.33)
758    parser.add_argument("--random_angle_range", help="Generate a random angle within [-random_angle_range, random_angle_range] to offset angles",
759                        type=float, default=20)
760    parser.add_argument("--clumping_factor", type=float, default=1.0)
761    parser.add_argument("--leaf_distance_factor_to_branch_root", type=float, default=0.3)
762
763    parser.add_argument("--out_obj_path", help="", type=str, required=True)
764    args = parser.parse_args()
765
766    cg = CrownGenerator()
767    cg.crown_shape = args.crown_shape
768    cg.crown_height = args.crown_height
769    cg.crown_diameter_SN = args.crown_diameter_sn
770    cg.crown_diameter_EW = args.crown_diameter_ew
771    if args.trunk_height > 0:
772        cg.trunk_height = args.trunk_height
773        cg.has_trunk = True
774    else:
775        cg.has_trunk = False
776    cg.dbh = args.dbh
777    if cg.dbh <= 0:
778        cg.has_trunk = False
779    cg.leaf_angle_dist = args.lad
780    cg.leaf_volume_density = args.lvd
781    cg.leaf_shape = args.leaf_shape
782    cg.leaf_num_triangles = args.polygon_sides
783    cg.single_leaf_area = args.single_leaf_area
784    cg.sig_lower = args.lower_sigma
785    cg.sig_upper = args.upper_sigma
786    cg.gamma_shape = args.gamma_shape
787    cg.generate_crown_only = args.generate_crown_only
788
789    cg.has_branches = args.has_branches
790    cg.branch_vertical_angle = args.branch_vertical_angle
791    cg.vertical_distance_between_branch = args.vertical_distance_between_branch
792    cg.num_branches_each_height = args.num_branches_each_height
793    cg.branch_radius_scale = args.branch_radius_scale
794    cg.random_angle_range = args.random_angle_range
795    if cg.branch_radius_scale <= 0 or cg.branch_vertical_angle <= 0 or cg.branch_vertical_angle >= 180:
796        cg.has_branches = False
797    cg.clumping_factor = args.clumping_factor
798    cg.leaf_distance_factor_to_branch_root = args.leaf_distance_factor_to_branch_root
799    cg.generate_crown(args.out_obj_path)
800
801    # leaf_nums = [1304,1467,1630,1793,1956,2119,2282,2445,2608,2771,2934,3097,3259,3422,3585]
802    # for i in range(len(leaf_nums)):
803    #     leaf_num = leaf_nums[i]
804    #     cg = CrownGenerator()
805    #     cg.crown_shape = CrownShape.CUBE
806    #     cg.crown_height = 0.4
807    #     cg.crown_diameter_SN = 4
808    #     cg.crown_diameter_EW = 4
809    #     cg.trunk_height = 5 + i*0.4
810    #     cg.has_trunk = True
811    #     cg.dbh = 0.4
812    #     cg.leaf_angle_dist = LAD.SPHERICAL
813    #     cg.num_of_leaves = leaf_num
814    #     cg.leaf_shape = LeafShape.DISK
815    #     cg.leaf_num_triangles = 6
816    #     cg.single_leaf_area = 0.00196350
817    #     cg.generate_crown(r"d:\TMP\facet_layer_"+"%03d"%(i+1)+".obj")