Update app.py

app.py

@@ -1,8 +1,10 @@
 import gradio as gr
 import spaces
 from gradio_litmodel3d import LitModel3D
+
 import os
 import shutil
+import trimesh  # New import
 os.environ['SPCONV_ALGO'] = 'native'
 from typing import *
 import torch
@@ -13,303 +15,157 @@ from PIL import Image
 from trellis.pipelines import TrellisImageTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
+from scipy.spatial import ConvexHull  # New import
 
-MAX_SEED = np.iinfo(np.int32).max
-TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
-os.makedirs(TMP_DIR, exist_ok=True)
-
-def start_session(req: gr.Request):
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    os.makedirs(user_dir, exist_ok=True)
-    
-def end_session(req: gr.Request):
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    shutil.rmtree(user_dir)
+# [Previous imports and constants remain the same...]
 
-def preprocess_image(image: Image.Image) -> Image.Image:
-    """
-    Preprocess the input image.
-    Args:
-        image (Image.Image): The input image.
-    Returns:
-        Image.Image: The preprocessed image.
+def optimize_building_mesh(mesh, angle_threshold=15, planar_threshold=0.02):
     """
-    processed_image = pipeline.preprocess_image(image)
-    return processed_image
-
-def preprocess_images(images: List[Tuple[Image.Image, str]]) -> List[Image.Image]:
+    Optimize a building mesh by preserving architectural features while reducing complexity.
     """
-    Preprocess a list of input images.
+    # Convert vertices to numpy array for processing
+    vertices = np.array(mesh.vertices)
+    faces = np.array(mesh.faces)
+    
+    # 1. Detect planar surfaces
+    normals = mesh.face_normals
+    planar_groups = []
+    processed = set()
     
-    Args:
-        images (List[Tuple[Image.Image, str]]): The input images.
+    for i in range(len(faces)):
+        if i in processed:
+            continue
         
-    Returns:
-        List[Image.Image]: The preprocessed images.
-    """
-    images = [image[0] for image in images]
-    processed_images = [pipeline.preprocess_image(image) for image in images]
-    return processed_images
-
-def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
-    return {
-        'gaussian': {
-            **gs.init_params,
-            '_xyz': gs._xyz.cpu().numpy(),
-            '_features_dc': gs._features_dc.cpu().numpy(),
-            '_scaling': gs._scaling.cpu().numpy(),
-            '_rotation': gs._rotation.cpu().numpy(),
-            '_opacity': gs._opacity.cpu().numpy(),
-        },
-        'mesh': {
-            'vertices': mesh.vertices.cpu().numpy(),
-            'faces': mesh.faces.cpu().numpy(),
-        },
-    }
+        # Find connected faces with similar normals
+        similar_faces = {i}
+        stack = [i]
+        while stack:
+            current = stack.pop()
+            neighbors = mesh.face_adjacency[mesh.face_adjacency[:,0] == current][:,1]
+            for n in neighbors:
+                if n not in processed:
+                    angle = np.arccos(np.dot(normals[current], normals[n])) * 180 / np.pi
+                    if angle < angle_threshold:
+                        similar_faces.add(n)
+                        stack.append(n)
+                        processed.add(n)
+        
+        if len(similar_faces) > 0:
+            planar_groups.append(list(similar_faces))
     
-def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
-    gs = Gaussian(
-        aabb=state['gaussian']['aabb'],
-        sh_degree=state['gaussian']['sh_degree'],
-        mininum_kernel_size=state['gaussian']['mininum_kernel_size'],
-        scaling_bias=state['gaussian']['scaling_bias'],
-        opacity_bias=state['gaussian']['opacity_bias'],
-        scaling_activation=state['gaussian']['scaling_activation'],
-    )
-    gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda')
-    gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda')
-    gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
-    gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
-    gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
+    # 2. Simplify each planar group while preserving edges
+    new_vertices = []
+    new_faces = []
+    vertex_map = {}
+    
+    for group in planar_groups:
+        # Get vertices for this group
+        group_faces = faces[group]
+        group_verts = vertices[np.unique(group_faces)]
+        
+        # Find best fitting plane
+        centroid = np.mean(group_verts, axis=0)
+        _, _, vh = np.linalg.svd(group_verts - centroid)
+        normal = vh[2]
+        
+        # Project vertices to plane and simplify
+        projected = group_verts - np.dot(group_verts - centroid, normal)[:, np.newaxis] * normal
+        
+        # Create simplified convex hull for this section
+        hull = ConvexHull(projected[:,:2])
+        hull_vertices = projected[hull.vertices]
+        
+        # Add to new mesh
+        start_idx = len(new_vertices)
+        new_vertices.extend(hull_vertices)
+        
+        # Triangulate the hull
+        for i in range(1, len(hull_vertices) - 1):
+            new_faces.append([start_idx, start_idx + i, start_idx + i + 1])
     
-    mesh = edict(
-        vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
-        faces=torch.tensor(state['mesh']['faces'], device='cuda'),
+    # 3. Create new optimized mesh
+    optimized_mesh = trimesh.Trimesh(
+        vertices=np.array(new_vertices),
+        faces=np.array(new_faces)
     )
     
-    return gs, mesh
-
-def get_seed(randomize_seed: bool, seed: int) -> int:
-    """
-    Get the random seed.
-    """
-    return np.random.randint(0, MAX_SEED) if randomize_seed else seed
-
-@spaces.GPU
-def image_to_3d(
-    image: Image.Image,
-    multiimages: List[Tuple[Image.Image, str]],
-    is_multiimage: bool,
-    seed: int,
-    ss_guidance_strength: float,
-    ss_sampling_steps: int,
-    slat_guidance_strength: float,
-    slat_sampling_steps: int,
-    multiimage_algo: Literal["multidiffusion", "stochastic"],
-    req: gr.Request,
-) -> Tuple[dict, str]:
-    """
-    Convert an image to a 3D model.
-    Args:
-        image (Image.Image): The input image.
-        multiimages (List[Tuple[Image.Image, str]]): The input images in multi-image mode.
-        is_multiimage (bool): Whether is in multi-image mode.
-        seed (int): The random seed.
-        ss_guidance_strength (float): The guidance strength for sparse structure generation.
-        ss_sampling_steps (int): The number of sampling steps for sparse structure generation.
-        slat_guidance_strength (float): The guidance strength for structured latent generation.
-        slat_sampling_steps (int): The number of sampling steps for structured latent generation.
-        multiimage_algo (Literal["multidiffusion", "stochastic"]): The algorithm for multi-image generation.
-    Returns:
-        dict: The information of the generated 3D model.
-        str: The path to the video of the 3D model.
-    """
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    if not is_multiimage:
-        outputs = pipeline.run(
-            image,
-            seed=seed,
-            formats=["gaussian", "mesh"],
-            preprocess_image=False,
-            sparse_structure_sampler_params={
-                "steps": ss_sampling_steps,
-                "cfg_strength": ss_guidance_strength,
-            },
-            slat_sampler_params={
-                "steps": slat_sampling_steps,
-                "cfg_strength": slat_guidance_strength,
-            },
-        )
-    else:
-        outputs = pipeline.run_multi_image(
-            [image[0] for image in multiimages],
-            seed=seed,
-            formats=["gaussian", "mesh"],
-            preprocess_image=False,
-            sparse_structure_sampler_params={
-                "steps": ss_sampling_steps,
-                "cfg_strength": ss_guidance_strength,
-            },
-            slat_sampler_params={
-                "steps": slat_sampling_steps,
-                "cfg_strength": slat_guidance_strength,
-            },
-            mode=multiimage_algo,
-        )
-    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
-    video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
-    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
-    video_path = os.path.join(user_dir, 'sample.mp4')
-    imageio.mimsave(video_path, video, fps=15)
-    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
-    torch.cuda.empty_cache()
-    return state, video_path
+    return optimized_mesh
 
+# Modify the existing extract_glb function
 @spaces.GPU(duration=90)
 def extract_glb(
     state: dict,
     mesh_simplify: float,
     texture_size: int,
+    is_building: bool,  # New parameter
+    angle_threshold: float,  # New parameter
+    planar_threshold: float,  # New parameter
     req: gr.Request,
 ) -> Tuple[str, str]:
     """
-    Extract a GLB file from the 3D model.
-    Args:
-        state (dict): The state of the generated 3D model.
-        mesh_simplify (float): The mesh simplification factor.
-        texture_size (int): The texture resolution.
-    Returns:
-        str: The path to the extracted GLB file.
+    Extract a GLB file from the 3D model with optional building optimization.
     """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     gs, mesh = unpack_state(state)
+    
+    if is_building:
+        # Convert to trimesh for optimization
+        trimesh_mesh = trimesh.Trimesh(
+            vertices=mesh.vertices.cpu().numpy(),
+            faces=mesh.faces.cpu().numpy()
+        )
+        
+        # Apply building-specific optimization
+        optimized_mesh = optimize_building_mesh(
+            trimesh_mesh,
+            angle_threshold=angle_threshold,
+            planar_threshold=planar_threshold
+        )
+        
+        # Convert back to original format
+        mesh.vertices = torch.tensor(optimized_mesh.vertices, device='cuda')
+        mesh.faces = torch.tensor(optimized_mesh.faces, device='cuda')
+    
     glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
     glb_path = os.path.join(user_dir, 'sample.glb')
     glb.export(glb_path)
     torch.cuda.empty_cache()
     return glb_path, glb_path
 
-@spaces.GPU
-def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
-    """
-    Extract a Gaussian file from the 3D model.
-    Args:
-        state (dict): The state of the generated 3D model.
-    Returns:
-        str: The path to the extracted Gaussian file.
-    """
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    gs, _ = unpack_state(state)
-    gaussian_path = os.path.join(user_dir, 'sample.ply')
-    gs.save_ply(gaussian_path)
-    torch.cuda.empty_cache()
-    return gaussian_path, gaussian_path
-
-with gr.Blocks(theme=gr.themes.Default(), delete_cache=(600, 600)) as demo:
-    with gr.Row():
-        with gr.Column():
-            with gr.Tabs() as input_tabs:
-                with gr.Tab(label="Single Image", id=0) as single_image_input_tab:
-                    image_prompt = gr.Image(label="Image Prompt", format="png", image_mode="RGBA", type="pil", height=300)
-                with gr.Tab(label="Multiple Images", id=1) as multiimage_input_tab:
-                    multiimage_prompt = gr.Gallery(label="Image Prompt", format="png", type="pil", height=300, columns=3)
-            
-            with gr.Accordion(label="Generation Settings", open=False):
-                seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
-                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-                with gr.Row():
-                    ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
-                    ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
-                with gr.Row():
-                    slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1)
-                    slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
-                multiimage_algo = gr.Radio(["stochastic", "multidiffusion"], label="Multi-image Algorithm", value="stochastic")
-
-            generate_btn = gr.Button("Generate", variant="primary")
-            
-            with gr.Accordion(label="GLB Extraction Settings", open=False):
-                mesh_simplify = gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01)
-                texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
-            
-            with gr.Row():
-                extract_glb_btn = gr.Button("Extract GLB", interactive=False)
-                extract_gs_btn = gr.Button("Extract Gaussian", interactive=False)
-
-        with gr.Column():
-            video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
-            model_output = LitModel3D(label="Extracted GLB/Gaussian", exposure=10.0, height=300)
-            
-            with gr.Row():
-                download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
-                download_gs = gr.DownloadButton(label="Download Gaussian", interactive=False)  
+# Modify the main UI code section
+with gr.Blocks(delete_cache=(600, 600)) as demo:
+    # [Previous UI code remains the same until GLB Extraction Settings...]
     
-    is_multiimage = gr.State(False)
-    output_buf = gr.State()
-
-    # Handlers
-    demo.load(start_session)
-    demo.unload(end_session)
+    with gr.Accordion(label="GLB Extraction Settings", open=False):
+        mesh_simplify = gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01)
+        texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
+        # Add new building optimization controls
+        with gr.Row():
+            is_building = gr.Checkbox(label="Enable Building Optimization", value=False)
+        with gr.Column(visible=False) as building_settings:
+            angle_threshold = gr.Slider(5, 45, label="Edge Angle Threshold", value=15, step=1)
+            planar_threshold = gr.Slider(0.01, 0.1, label="Planar Surface Threshold", value=0.02, step=0.01)
     
-    single_image_input_tab.select(
-        lambda: False,
-        outputs=[is_multiimage]
-    )
-    multiimage_input_tab.select(
-        lambda: True,
-        outputs=[is_multiimage]
-    )
+    # [Rest of the UI code remains the same until the event handlers...]
     
-    image_prompt.upload(
-        preprocess_image,
-        inputs=[image_prompt],
-        outputs=[image_prompt],
+    # Add visibility toggle for building settings
+    is_building.change(
+        lambda x: gr.Column.update(visible=x),
+        inputs=[is_building],
+        outputs=[building_settings]
     )
-    multiimage_prompt.upload(
-        preprocess_images,
-        inputs=[multiimage_prompt],
-        outputs=[multiimage_prompt],
-    )
-
-    generate_btn.click(
-        get_seed,
-        inputs=[randomize_seed, seed],
-        outputs=[seed],
-    ).then(
-        image_to_3d,
-        inputs=[image_prompt, multiimage_prompt, is_multiimage, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps, multiimage_algo],
-        outputs=[output_buf, video_output],
-    ).then(
-        lambda: tuple([gr.Button(interactive=True), gr.Button(interactive=True)]),
-        outputs=[extract_glb_btn, extract_gs_btn],
-    )
-
-    video_output.clear(
-        lambda: tuple([gr.Button(interactive=False), gr.Button(interactive=False)]),
-        outputs=[extract_glb_btn, extract_gs_btn],
-    )
-
+    
+    # Modify the extract_glb button click handler
     extract_glb_btn.click(
         extract_glb,
-        inputs=[output_buf, mesh_simplify, texture_size],
+        inputs=[output_buf, mesh_simplify, texture_size, is_building, angle_threshold, planar_threshold],
         outputs=[model_output, download_glb],
     ).then(
         lambda: gr.Button(interactive=True),
         outputs=[download_glb],
     )
     
-    extract_gs_btn.click(
-        extract_gaussian,
-        inputs=[output_buf],
-        outputs=[model_output, download_gs],
-    ).then(
-        lambda: gr.Button(interactive=True),
-        outputs=[download_gs],
-    )
-
-    model_output.clear(
-        lambda: gr.Button(interactive=False),
-        outputs=[download_glb],
-    )
+    # [Rest of the code remains the same...]
 
 # Launch the Gradio app
 if __name__ == "__main__":
@@ -319,4 +175,4 @@ if __name__ == "__main__":
         pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))    # Preload rembg
     except:
         pass
-    demo.launch()
+    demo.launch()