Update app.py

app.py

@@ -286,11 +286,6 @@ class GPUSatelliteModelGenerator:
         
         # Detect tree areas and generate tree geometry
         tree_mask = cp.all(texture_img_gpu == self.colors['dark_green'], axis=2)
-        
-        # 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
@@ -303,9 +298,9 @@ class GPUSatelliteModelGenerator:
         faces = self.generate_faces_gpu(height, width)
         uvs = self.generate_uvs_gpu(vertices, width, height)
         
-        # Create textured mesh
+        # Create textured mesh using the original texture image
         return self.create_textured_mesh(vertices, faces, uvs, texture_img)
-        
+            
 
     @staticmethod
     def generate_faces_gpu(height, width):
@@ -332,20 +327,32 @@ class GPUSatelliteModelGenerator:
     @staticmethod
     def create_textured_mesh(vertices, faces, uvs, texture_img):
         """Create textured mesh with proper color conversion"""
-        if len(texture_img.shape) == 3 and texture_img.shape[2] == 4:
-            texture_img = cv2.cvtColor(texture_img, cv2.COLOR_BGRA2RGB)
-        elif len(texture_img.shape) == 3:
-            texture_img = cv2.cvtColor(texture_img, cv2.COLOR_BGR2RGB)
-        
-        return trimesh.Trimesh(
-            vertices=vertices.get(),
-            faces=faces.get(),
+        # Ensure we're working with the original texture image
+        if isinstance(texture_img, cp.ndarray):
+            texture_img = texture_img.get()
+            
+        # Convert texture image to RGB format for PIL
+        if len(texture_img.shape) == 3:
+            if texture_img.shape[2] == 4:  # BGRA
+                texture_img = cv2.cvtColor(texture_img, cv2.COLOR_BGRA2RGB)
+            else:  # BGR
+                texture_img = cv2.cvtColor(texture_img, cv2.COLOR_BGR2RGB)
+        
+        # Create PIL Image from the texture
+        texture_pil = Image.fromarray(texture_img)
+        
+        # Create the mesh with texture
+        mesh = trimesh.Trimesh(
+            vertices=vertices.get() if isinstance(vertices, cp.ndarray) else vertices,
+            faces=faces.get() if isinstance(faces, cp.ndarray) else faces,
             visual=trimesh.visual.TextureVisuals(
-                uv=uvs.get(),
-                image=Image.fromarray(texture_img)
+                uv=uvs.get() if isinstance(uvs, cp.ndarray) else uvs,
+                image=texture_pil
             )
         )
         
+        return mesh
+        
 def generate_and_process_map(prompt: str) -> tuple[str | None, np.ndarray | None]:
     """Generate satellite image from prompt and convert to 3D model using GPU acceleration"""
     try: