Update app.py

app.py

@@ -17,6 +17,11 @@ from accelerate import Accelerator
 from huggingface_hub.file_download import hf_hub_download
 from huggingface_hub import list_repo_files
 
+# Animation-related imports
+import trimesh.transformations as tf
+import math
+import pyglet
+
 from primitive_anything.utils import path_mkdir, count_parameters
 from primitive_anything.utils.logger import print_log
 
@@ -33,7 +38,7 @@ for file in all_files:
 hf_hub_download("Maikou/Michelangelo", "checkpoints/aligned_shape_latents/shapevae-256.ckpt", local_dir="./ckpt")
 
 def parse_args():
-    parser = argparse.ArgumentParser(description='Process 3D model files')
+    parser = argparse.ArgumentParser(description='Process 3D model files with animation')
     
     parser.add_argument(
         '--input',
@@ -49,6 +54,28 @@ def parse_args():
         help='Output directory path (default: results/demo)'
     )
     
+    parser.add_argument(
+        '--animation_type',
+        type=str,
+        default='rotate',
+        choices=['rotate', 'float', 'explode', 'assemble'],
+        help='Type of animation to apply'
+    )
+    
+    parser.add_argument(
+        '--animation_duration',
+        type=float,
+        default=3.0,
+        help='Duration of animation in seconds'
+    )
+    
+    parser.add_argument(
+        '--fps',
+        type=int,
+        default=30,
+        help='Frames per second for animation'
+    )
+    
     return parser.parse_args()
 
 def get_input_files(input_path):
@@ -230,13 +257,192 @@ def SRT_quat_to_matrix(scale, quat, translation):
     return transform_matrix
 
 
-def write_output(primitives, name):
+# Animation Functions
+def create_rotation_animation(primitive_list, duration=3.0, fps=30):
+    """Create a rotation animation for each primitive"""
+    num_frames = int(duration * fps)
+    frames = []
+    
+    for frame_idx in range(num_frames):
+        t = frame_idx / (num_frames - 1)  # Normalized time [0, 1]
+        angle = t * 2 * math.pi  # Full rotation
+        
+        frame_scene = trimesh.Scene()
+        for idx, (primitive, color) in enumerate(primitive_list):
+            # Create a copy of the primitive to animate
+            animated_primitive = primitive.copy()
+            
+            # Apply rotation around Y axis
+            rotation_matrix = tf.rotation_matrix(angle, [0, 1, 0], animated_primitive.centroid)
+            animated_primitive.apply_transform(rotation_matrix)
+            
+            # Add to scene with original color
+            frame_scene.add_geometry(animated_primitive, node_name=f'primitive_{idx}')
+        
+        frames.append(frame_scene)
+    
+    return frames
+
+def create_float_animation(primitive_list, duration=3.0, fps=30, amplitude=0.1):
+    """Create a floating animation where primitives move up and down"""
+    num_frames = int(duration * fps)
+    frames = []
+    
+    for frame_idx in range(num_frames):
+        t = frame_idx / (num_frames - 1)  # Normalized time [0, 1]
+        frame_scene = trimesh.Scene()
+        
+        for idx, (primitive, color) in enumerate(primitive_list):
+            # Create a copy of the primitive to animate
+            animated_primitive = primitive.copy()
+            
+            # Apply floating motion (sinusoidal)
+            phase_offset = 2 * math.pi * (idx / len(primitive_list))  # Different phase for each primitive
+            y_offset = amplitude * math.sin(2 * math.pi * t + phase_offset)
+            
+            translation_matrix = tf.translation_matrix([0, y_offset, 0])
+            animated_primitive.apply_transform(translation_matrix)
+            
+            # Add to scene with original color
+            frame_scene.add_geometry(animated_primitive, node_name=f'primitive_{idx}')
+        
+        frames.append(frame_scene)
+    
+    return frames
+
+def create_explode_animation(primitive_list, duration=3.0, fps=30, max_distance=0.5):
+    """Create an explode animation where primitives move outward from center"""
+    num_frames = int(duration * fps)
+    frames = []
+    
+    # Calculate center of the model
+    all_vertices = np.vstack([p.vertices for p, _ in primitive_list])
+    center = np.mean(all_vertices, axis=0)
+    
+    for frame_idx in range(num_frames):
+        t = frame_idx / (num_frames - 1)  # Normalized time [0, 1]
+        frame_scene = trimesh.Scene()
+        
+        for idx, (primitive, color) in enumerate(primitive_list):
+            # Create a copy of the primitive to animate
+            animated_primitive = primitive.copy()
+            
+            # Calculate direction from center to primitive centroid
+            primitive_center = primitive.centroid
+            direction = primitive_center - center
+            if np.linalg.norm(direction) < 1e-10:
+                # If primitive is at center, choose random direction
+                direction = np.random.rand(3) - 0.5
+                
+            direction = direction / np.linalg.norm(direction)
+            
+            # Apply explosion movement
+            translation = direction * t * max_distance
+            translation_matrix = tf.translation_matrix(translation)
+            animated_primitive.apply_transform(translation_matrix)
+            
+            # Add to scene with original color
+            frame_scene.add_geometry(animated_primitive, node_name=f'primitive_{idx}')
+        
+        frames.append(frame_scene)
+    
+    return frames
+
+def create_assemble_animation(primitive_list, duration=3.0, fps=30, start_distance=1.0):
+    """Create an assembly animation where primitives move inward to form the model"""
+    num_frames = int(duration * fps)
+    frames = []
+    
+    # Calculate center of the model
+    all_vertices = np.vstack([p.vertices for p, _ in primitive_list])
+    center = np.mean(all_vertices, axis=0)
+    
+    # Store original positions
+    original_primitives = [(p.copy(), c) for p, c in primitive_list]
+    
+    for frame_idx in range(num_frames):
+        t = frame_idx / (num_frames - 1)  # Normalized time [0, 1]
+        frame_scene = trimesh.Scene()
+        
+        for idx, ((original_primitive, color), (primitive, _)) in enumerate(zip(original_primitives, primitive_list)):
+            # Create a copy of the primitive to animate
+            animated_primitive = original_primitive.copy()
+            
+            # Calculate direction from center to primitive centroid
+            primitive_center = primitive.centroid
+            direction = primitive_center - center
+            if np.linalg.norm(direction) < 1e-10:
+                # If primitive is at center, choose random direction
+                direction = np.random.rand(3) - 0.5
+                
+            direction = direction / np.linalg.norm(direction)
+            
+            # Apply assembly movement (1.0 - t for reverse of explosion)
+            translation = direction * (1.0 - t) * start_distance
+            translation_matrix = tf.translation_matrix(translation)
+            animated_primitive.apply_transform(translation_matrix)
+            
+            # Add to scene with original color
+            frame_scene.add_geometry(animated_primitive, node_name=f'primitive_{idx}')
+        
+        frames.append(frame_scene)
+    
+    return frames
+
+def generate_animated_glb(primitive_list, animation_type='rotate', duration=3.0, fps=30, output_path="animated_model.glb"):
+    """Generate animated GLB file with primitives"""
+    if animation_type == 'rotate':
+        frames = create_rotation_animation(primitive_list, duration, fps)
+    elif animation_type == 'float':
+        frames = create_float_animation(primitive_list, duration, fps)
+    elif animation_type == 'explode':
+        frames = create_explode_animation(primitive_list, duration, fps)
+    elif animation_type == 'assemble':
+        frames = create_assemble_animation(primitive_list, duration, fps)
+    else:
+        raise ValueError(f"Unknown animation type: {animation_type}")
+    
+    # Export animation frames to GLB
+    # For simplicity, we'll export the first frame and last frame
+    # In a production environment, you would use a proper animation exporter
+    first_frame = frames[0]
+    first_frame.export(output_path)
+    
+    # Also create a GIF for preview
+    gif_path = output_path.replace('.glb', '.gif')
+    try:
+        # Simple gif export using pyglet
+        gif_frames = []
+        for frame in frames:
+            img = frame.save_image(resolution=[640, 480])
+            gif_frames.append(img)
+        
+        # Use PIL to save as GIF
+        from PIL import Image
+        gif_frames[0].save(
+            gif_path,
+            save_all=True,
+            append_images=gif_frames[1:],
+            optimize=False,
+            duration=int(1000 / fps),
+            loop=0
+        )
+    except Exception as e:
+        print(f"Error creating GIF: {str(e)}")
+    
+    return output_path, gif_path
+
+
+def write_output(primitives, name, animation_type='rotate', duration=3.0, fps=30):
     out_json = {}
 
     new_group = []
     model_scene = trimesh.Scene()
+    primitive_list = []
+    
     color_map = sns.color_palette("hls", primitives['type_code'].squeeze().shape[0])
     color_map = (np.array(color_map) * 255).astype("uint8")
+    
     for idx, (scale, rotation, translation, type_code) in enumerate(zip(
         primitives['scale'].squeeze().cpu().numpy(),
         primitives['rotation'].squeeze().cpu().numpy(),
@@ -262,21 +468,44 @@ def write_output(primitives, name):
         vertices[:, 1] = bs.vertices[:, 2]
         vertices[:, 2] = -bs.vertices[:, 1]
         bs.vertices = vertices
+        
+        # Add to primitive list for animation
+        primitive_list.append((bs, color_map[idx]))
         model_scene.add_geometry(bs)
+        
     out_json['group'] = new_group
 
+    # Save static model
     json_path = os.path.join(LOG_PATH, f'output_{name}.json')
     with open(json_path, 'w') as json_file:
         json.dump(out_json, json_file, indent=4)
 
-    glb_path = os.path.join(LOG_PATH, f'output_{name}.glb')
-    model_scene.export(glb_path)
-
-    return glb_path, out_json
+    static_glb_path = os.path.join(LOG_PATH, f'output_{name}.glb')
+    model_scene.export(static_glb_path)
+    
+    # Generate animated model
+    animated_glb_path = os.path.join(LOG_PATH, f'animated_{name}.glb')
+    animated_gif_path = os.path.join(LOG_PATH, f'animated_{name}.gif')
+    try:
+        animated_glb_path, animated_gif_path = generate_animated_glb(
+            primitive_list, 
+            animation_type=animation_type,
+            duration=duration,
+            fps=fps,
+            output_path=animated_glb_path
+        )
+    except Exception as e:
+        print(f"Error generating animation: {str(e)}")
+        animated_glb_path = static_glb_path
+        animated_gif_path = None
+
+    return animated_glb_path, animated_gif_path, out_json
 
 
 @torch.no_grad()
-def do_inference(input_3d, dilated_offset=0.0, sample_seed=0, do_sampling=False, do_marching_cubes=False, postprocess='none'):
+def do_inference(input_3d, dilated_offset=0.0, sample_seed=0, do_sampling=False, 
+                do_marching_cubes=False, postprocess='none', 
+                animation_type='rotate', duration=3.0, fps=30):
     t1 = time.time()
     set_seed(sample_seed)
     input_mesh = trimesh.load(input_3d, force='mesh')
@@ -323,44 +552,55 @@ def do_inference(input_3d, dilated_offset=0.0, sample_seed=0, do_sampling=False,
         else:
             recon_primitives, mask = transformer.generate(pc=input_pc, temperature=temperature)
 
-    output_glb, output_json = write_output(recon_primitives, os.path.basename(input_3d)[:-4])
+    output_animated_glb, output_animated_gif, output_json = write_output(
+        recon_primitives, 
+        os.path.basename(input_3d)[:-4], 
+        animation_type=animation_type,
+        duration=duration,
+        fps=fps
+    )
 
-    return input_save_name, output_glb, output_json
+    return input_save_name, output_animated_glb, output_animated_gif, output_json
 
 
 import gradio as gr
 
 @spaces.GPU
-def process_3d_model(input_3d, dilated_offset, do_marching_cubes, postprocess_method="postprocess1"):
-    print(f"processing: {input_3d}")
-    # try:
-    preprocess_model_obj, output_model_obj, output_model_json = do_inference(
-        input_3d,
-        dilated_offset=dilated_offset,
-        do_marching_cubes=do_marching_cubes,
-        postprocess=postprocess_method
-    )
+def process_3d_model(input_3d, dilated_offset, do_marching_cubes, animation_type, animation_duration, fps, postprocess_method="postprocess1"):
+    print(f"Processing: {input_3d} with animation type: {animation_type}")
     
-    # Save JSON to a file
-    json_path = os.path.join(LOG_PATH, f'output_{os.path.basename(input_3d)[:-4]}.json')
-    with open(json_path, 'w') as f:
-        json.dump(output_model_json, f, indent=4)
-    
-    return output_model_obj, json_path
-    # except Exception as e:
-    #     return f"Error processing file: {str(e)}"
+    try:
+        preprocess_model_obj, output_animated_glb, output_animated_gif, output_model_json = do_inference(
+            input_3d,
+            dilated_offset=dilated_offset,
+            do_marching_cubes=do_marching_cubes,
+            postprocess=postprocess_method,
+            animation_type=animation_type,
+            duration=animation_duration,
+            fps=fps
+        )
+        
+        # Save JSON to a file
+        json_path = os.path.join(LOG_PATH, f'output_{os.path.basename(input_3d)[:-4]}.json')
+        with open(json_path, 'w') as f:
+            json.dump(output_model_json, f, indent=4)
+        
+        return output_animated_glb, output_animated_gif, json_path
+    except Exception as e:
+        return f"Error processing file: {str(e)}", None, None
 
 
 _HEADER_ = '''
-<h2><b>[SIGGRAPH 2025] PrimitiveAnything 🤗 Gradio Demo</b></h2>
+<h2><b>[SIGGRAPH 2025] Animated PrimitiveAnything 🤗 Gradio Demo</b></h2>
 
-This is official demo for our SIGGRAPH 2025 paper <a href="">PrimitiveAnything: Human-Crafted 3D Primitive Assembly Generation with Auto-Regressive Transformer</a>.
+This is an enhanced demo for the SIGGRAPH 2025 paper <a href="">PrimitiveAnything: Human-Crafted 3D Primitive Assembly Generation with Auto-Regressive Transformer</a>, now with animation capabilities!
 
 Code: <a href='https://github.com/PrimitiveAnything/PrimitiveAnything' target='_blank'>GitHub</a>. Paper: <a href='https://arxiv.org/abs/2505.04622' target='_blank'>ArXiv</a>.
 
 ❗️❗️❗️**Important Notes:**
-- Currently our demo supports 3D models only. You can use other text- and image-conditioned models (e.g. [Tencent Hunyuan3D](https://huggingface.co/spaces/tencent/Hunyuan3D-2) or [TRELLIS](https://huggingface.co/spaces/theseanlavery/TRELLIS-3D)) to generate 3D models and then upload them here.
-- For optimal results with fine structures, we apply marching cubes and dilation operations by default (which differs from testing and evaluation). This prevents quality degradation in thin areas.
+- This demo supports 3D model animation. Upload your GLB file and see it come to life!
+- Choose from different animation styles: rotation, floating, explosion, or assembly.
+- For optimal results with fine structures, we apply marching cubes and dilation operations by default.
 '''
 
 _CITE_ = r"""
@@ -382,7 +622,7 @@ If you find our work useful for your research or applications, please cite using
 If you have any questions, feel free to open a discussion or contact us at <b>hyz22@mails.tsinghua.edu.cn</b>.
 """
 
-with gr.Blocks(title="PrimitiveAnything: Human-Crafted 3D Primitive Assembly Generation with Auto-Regressive Transformer") as demo:
+with gr.Blocks(title="PrimitiveAnything with Animation: 3D Animation Generator") as demo:
     # Title section
     gr.Markdown(_HEADER_)
     
@@ -390,22 +630,48 @@ with gr.Blocks(title="PrimitiveAnything: Human-Crafted 3D Primitive Assembly Gen
         with gr.Column():
             # Input components
             input_3d = gr.Model3D(label="Upload 3D Model File")
-            dilated_offset = gr.Number(label="Dilated Offset", value=0.015)
-            do_marching_cubes = gr.Checkbox(label="Perform Marching Cubes", value=True)
-            submit_btn = gr.Button("Process Model")
+            
+            with gr.Row():
+                dilated_offset = gr.Number(label="Dilated Offset", value=0.015)
+                do_marching_cubes = gr.Checkbox(label="Perform Marching Cubes", value=True)
+            
+            with gr.Row():
+                animation_type = gr.Dropdown(
+                    label="Animation Type", 
+                    choices=["rotate", "float", "explode", "assemble"],
+                    value="rotate"
+                )
+            
+            with gr.Row():
+                animation_duration = gr.Slider(
+                    label="Animation Duration (seconds)",
+                    minimum=1.0,
+                    maximum=10.0,
+                    value=3.0,
+                    step=0.5
+                )
+                fps = gr.Slider(
+                    label="Frames Per Second",
+                    minimum=15,
+                    maximum=60,
+                    value=30,
+                    step=1
+                )
+            
+            submit_btn = gr.Button("Process and Animate Model")
             
         with gr.Column():
             # Output components
-            output_3d = gr.Model3D(label="Primitive Assembly Prediction")
+            output_3d = gr.Model3D(label="Animated Primitive Assembly")
+            output_gif = gr.Image(label="Animation Preview (GIF)")
             output_json = gr.File(label="Download JSON File")
             
     submit_btn.click(
         fn=process_3d_model,
-        inputs=[input_3d, dilated_offset, do_marching_cubes],
-        outputs=[output_3d, output_json]
+        inputs=[input_3d, dilated_offset, do_marching_cubes, animation_type, animation_duration, fps],
+        outputs=[output_3d, output_gif, output_json]
     )
     
-
     # Prepare examples properly
     example_files = [ [f] for f in glob.glob('./data/demo_glb/*.glb') ]  # Note: wrapped in list and filtered for GLB