Create utils.py

utils.py

@@ -0,0 +1,57 @@
+# timeforge/utils.py
+
+import tempfile
+import trimesh
+from trimesh.exchange.gltf import export_glb
+import numpy as np
+from PIL import Image
+
+def apply_gradient_color(mesh_text):
+    """
+    Apply a gradient color to the mesh vertices and return as GLB.
+    Args:
+        mesh_text (str): The input mesh in OBJ format as a string.
+    Returns:
+        str: Path to the GLB file with gradient colors applied.
+    """
+    # Load the mesh
+    temp_file =  tempfile.NamedTemporaryFile(suffix=f"", delete=False).name
+    with open(temp_file+".obj", "w") as f:
+        f.write(mesh_text)
+    mesh = trimesh.load_mesh(temp_file+".obj", file_type='obj')
+
+    # Get vertex coordinates
+    vertices = mesh.vertices
+    y_values = vertices[:, 1]  # Y-axis values
+
+    # Normalize Y values to range [0, 1] for color mapping
+    y_normalized = (y_values - y_values.min()) / (y_values.max() - y_values.min())
+
+    # Generate colors: Map normalized Y values to RGB gradient (e.g., blue to red)
+    colors = np.zeros((len(vertices), 4))  # RGBA
+    colors[:, 0] = y_normalized  # Red channel
+    colors[:, 2] = 1 - y_normalized  # Blue channel
+    colors[:, 3] = 1.0  # Alpha channel (fully opaque)
+
+    # Attach colors to mesh vertices
+    mesh.visual.vertex_colors = colors
+
+    # Export to GLB format
+    glb_path = temp_file+".glb"
+    with open(glb_path, "wb") as f:
+        f.write(export_glb(mesh))
+    
+    return glb_path
+
+def create_image_grid(images):
+    images = [Image.fromarray((img * 255).astype("uint8")) for img in images]
+
+    width, height = images[0].size
+    grid_img = Image.new("RGB", (2 * width, 2 * height))
+
+    grid_img.paste(images[0], (0, 0))
+    grid_img.paste(images[1], (width, 0))
+    grid_img.paste(images[2], (0, height))
+    grid_img.paste(images[3], (width, height))
+
+    return grid_img