Update app.py

app.py

@@ -2,99 +2,93 @@ import gradio as gr
 import torch
 import numpy as np
 from PIL import Image
-import os
 import tempfile
+import plyfile  # Add to requirements.txt
 
 class SimplePersistentCortex:
-    def __init__(self, num_gaussians=8000):
+    def __init__(self, num_gaussians=10000):
         self.num_gaussians = num_gaussians
-        # Initialize Gaussians centered in view
-        self.positions = torch.randn(num_gaussians, 3) * 0.3  # Tighter cluster
-        self.scales = torch.exp(torch.randn(num_gaussians, 3) * 0.5 + np.log(0.05))  # Log scale for positivity
-        self.colors = torch.sigmoid(torch.randn(num_gaussians, 3))  # RGB in [0,1]
-        self.opacities = torch.sigmoid(torch.randn(num_gaussians) * 2 + 1)  # Mostly opaque
+        # Initialize clustered Gaussians
+        self.positions = torch.randn(num_gaussians, 3) * 0.4
+        self.scales = torch.exp(torch.randn(num_gaussians, 3) * 0.6 - 1.0)  # Positive scales
+        self.colors_dc = torch.rand(num_gaussians, 3) * 0.5 + 0.25  # Base colors (SH degree 0)
+        self.opacities = torch.sigmoid(torch.randn(num_gaussians) * 1.5 + 2.0)  # Mostly visible
         self.rotations = torch.nn.functional.normalize(torch.randn(num_gaussians, 4), dim=-1)
 
-    def evolve(self, target_image: Image.Image, steps: int = 80):
-        # Resize and normalize target
-        target_np = np.array(target_image.resize((64, 64))) / 255.0
-        target_tensor = torch.tensor(target_np, dtype=torch.float32)
+    def evolve(self, target_image: Image.Image, steps: int = 100):
+        target_np = np.array(target_image.resize((128, 128))) / 255.0
+        target_tensor = torch.tensor(target_np, dtype=torch.float32).mean(dim=(0,1))  # Average color
 
-        # Simple attraction: move positions toward brighter areas (simulated depth)
         for _ in range(steps):
-            # Random small updates for "evolution"
-            self.positions += torch.randn_like(self.positions) * 0.008
+            # Attract positions to center with noise
+            self.positions += torch.randn_like(self.positions) * 0.01
+            self.positions *= 0.99  # Dampen spread
 
-            # Bias colors toward target average
-            avg_color = target_tensor.mean(dim=(0,1))
-            self.colors += (avg_color - self.colors.mean(dim=0)) * 0.01
-            self.colors = torch.clamp(self.colors, 0, 1)
+            # Adapt colors toward target
+            self.colors_dc += (target_tensor - self.colors_dc.mean(dim=0)) * 0.02
+            self.colors_dc = torch.clamp(self.colors_dc, 0, 1)
 
-            # Adjust scales for density
-            self.scales *= torch.exp(torch.randn_like(self.scales) * 0.02)
+            # Vary scales for density
+            self.scales *= torch.exp(torch.randn_like(self.scales) * 0.03)
 
-            # Increase opacity gradually
-            self.opacities = torch.clamp(self.opacities + 0.005, 0, 1)
+            # Boost opacity
+            self.opacities = torch.clamp(self.opacities + 0.01, 0, 1)
 
         return self
 
     def export_ply(self, path: str):
-        # Simple PLY export (ASCII vertex list; sufficient for demo)
-        with open(path, 'w') as f:
-            f.write("ply\nformat ascii 1.0\n")
-            f.write(f"element vertex {self.num_gaussians}\n")
-            f.write("property float x\nproperty float y\nproperty float z\n")
-            f.write("property float nx\nproperty float ny\nproperty float nz\n")  # Dummy normals
-            f.write("property float f_dc_0\nproperty float f_dc_1\nproperty float f_dc_2\n")  # SH DC
-            f.write("property float opacity\n")
-            f.write("property float scale_0\nproperty float scale_1\nproperty float scale_2\n")
-            f.write("property float rot_0\nproperty float rot_1\nproperty float rot_2\nproperty float rot_3\n")
-            f.write("end_header\n")
+        positions = self.positions.cpu().numpy()
+        scales = np.log(self.scales.cpu().numpy() + 1e-9)
+        opacities = np.arctanh(self.opacities.cpu().numpy().clip(1e-6, 1-1e-6))
+        rot = self.rotations.cpu().numpy()
+        sh_dc = self.colors_dc.cpu().numpy().reshape(-1, 3)  # Simplified SH
 
-            positions = self.positions.cpu().numpy()
-            scales = np.log(self.scales.cpu().numpy() + 1e-8)  # Log for PLY
-            colors = self.colors.cpu().numpy()
-            opacities = np.arctanh(self.opacities.cpu().numpy().clip(1e-6, 1-1e-6))  # Inverse sigmoid
-            rotations = self.rotations.cpu().numpy()
+        # Create vertex data
+        vertex_data = np.empty(self.num_gaussians, dtype=[
+            ('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
+            ('f_dc_0', 'f4'), ('f_dc_1', 'f4'), ('f_dc_2', 'f4'),
+            ('opacity', 'f4'),
+            ('scale_0', 'f4'), ('scale_1', 'f4'), ('scale_2', 'f4'),
+            ('rot_0', 'f4'), ('rot_1', 'f4'), ('rot_2', 'f4'), ('rot_3', 'f4')
+        ])
 
-            for i in range(self.num_gaussians):
-                x, y, z = positions[i]
-                nx, ny, nz = 0, 0, 0  # Dummy
-                r, g, b = colors[i]
-                op = opacities[i]
-                s0, s1, s2 = scales[i]
-                r0, r1, r2, r3 = rotations[i]
-                f.write(f"{x} {y} {z} {nx} {ny} {nz} {r} {g} {b} {op} {s0} {s1} {s2} {r0} {r1} {r2} {r3}\n")
+        vertex_data['x'], vertex_data['y'], vertex_data['z'] = positions[:, 0], positions[:, 1], positions[:, 2]
+        vertex_data['f_dc_0'], vertex_data['f_dc_1'], vertex_data['f_dc_2'] = sh_dc[:, 0], sh_dc[:, 1], sh_dc[:, 2]
+        vertex_data['opacity'] = opacities
+        vertex_data['scale_0'], vertex_data['scale_1'], vertex_data['scale_2'] = scales[:, 0], scales[:, 1], scales[:, 2]
+        vertex_data['rot_0'], vertex_data['rot_1'], vertex_data['rot_2'], vertex_data['rot_3'] = rot[:, 0], rot[:, 1], rot[:, 2], rot[:, 3]
+
+        el = plyfile.PlyElement.describe(vertex_data, 'vertex')
+        plyfile.PlyData([el]).write(path)
 
 def process(image: Image.Image, prompt: str = ""):
     if image is None:
-        return None, "Please upload an image to begin."
+        return None, "Error: Please upload an image."
     
     cortex = SimplePersistentCortex()
     cortex.evolve(image)
     
-    # Save to temporary PLY
     with tempfile.NamedTemporaryFile(delete=False, suffix=".ply") as tmp:
         cortex.export_ply(tmp.name)
         ply_path = tmp.name
     
-    status = "Persistent 3D Gaussian representation generated from input image. Interact with the viewer below (orbit, zoom)."
+    status = "Persistent 3D Gaussian splat generated. Use mouse to orbit/zoom in viewer."
     if prompt:
-        status += f" Prompt incorporated: {prompt}"
+        status += f" (Prompt: {prompt})"
     
     return ply_path, status
 
-with gr.Blocks(title="Persistent 3D Cortex Demo v2") as demo:
+with gr.Blocks(title="Persistent 3D Cortex Demo v3") as demo:
     gr.Markdown("# Persistent 3D Cortex Demo")
-    gr.Markdown("Upload an image (and optional prompt) to evolve a persistent 3D Gaussian splat representation.")
+    gr.Markdown("Upload an image to generate an evolving 3D Gaussian representation.")
     
     with gr.Row():
         img_input = gr.Image(type="pil", label="Input Image")
-        prompt_input = gr.Textbox(label="Prompt (optional)", placeholder="e.g., enhance depth")
+        prompt_input = gr.Textbox(label="Prompt (optional)")
     
-    btn = gr.Button("Generate Persistent 3D")
+    btn = gr.Button("Generate 3D")
     
-    model_output = gr.Model3D(label="Interactive 3D Viewer", clear_color=[0.1, 0.1, 0.1, 1.0])
+    model_output = gr.Model3D(label="Interactive Gaussian Splat Viewer", clear_color=[0.0, 0.0, 0.0, 1.0])
     text_output = gr.Textbox(label="Status")
     
     btn.click(process, inputs=[img_input, prompt_input], outputs=[model_output, text_output])