Update app.py

app.py

@@ -128,16 +128,21 @@ class GPUSatelliteModelGenerator:
         tree_positions = cp.where(tree_mask)
         num_trees = len(tree_positions[0])
         
+        # Return original vertices if no trees detected
+        if num_trees == 0:
+            return base_vertices
+        
         # Random height variation for trees
         tree_heights = cp.random.uniform(0.15, 0.25, num_trees)
         
         # Create vertex displacements for tree geometry
         tree_vertices = base_vertices.copy()
-        tree_vertices[tree_positions] += cp.stack([
-            cp.zeros(num_trees),  # x offset
-            tree_heights,         # y offset (height)
-            cp.zeros(num_trees)   # z offset
-        ], axis=1)
+        
+        # Get indices for tree positions
+        tree_indices = cp.ravel_multi_index(tree_positions, tree_mask.shape)
+        
+        # Add height offsets to tree positions
+        tree_vertices[tree_indices, 1] += tree_heights
         
         return tree_vertices
 
@@ -271,17 +276,22 @@ class GPUSatelliteModelGenerator:
     def generate_mesh_gpu(self, height_map, texture_img):
         """Generate optimized 3D mesh with tree geometry"""
         height_map_gpu = cp.asarray(height_map)
-        # Convert texture_img to CuPy array
         texture_img_gpu = cp.asarray(texture_img)
         height, width = height_map.shape
         
         # Generate base vertices
         x, z = cp.meshgrid(cp.arange(width), cp.arange(height))
         vertices = cp.stack([x, height_map_gpu * self.building_height, z], axis=-1)
+        vertices = vertices.reshape(-1, 3)
         
         # Detect tree areas and generate tree geometry
         tree_mask = cp.all(texture_img_gpu == self.colors['dark_green'], axis=2)
-        vertices = self.generate_tree_vertices(tree_mask, vertices.reshape(-1, 3))
+        
+        # Add some debug printing
+        print(f"Tree mask shape: {tree_mask.shape}")
+        print(f"Number of trees detected: {cp.sum(tree_mask)}")
+        
+        vertices = self.generate_tree_vertices(tree_mask, vertices)
         
         # Normalize coordinates
         scale = max(width, height)