Update app.py

app.py

@@ -2,95 +2,101 @@ import gradio as gr
 import torch
 import numpy as np
 from PIL import Image
-import gsplat  # Hugging Face provides gsplat in Spaces
+import os
+import tempfile
 
 class SimplePersistentCortex:
-    def __init__(self, num_gaussians=5000):
+    def __init__(self, num_gaussians=8000):
         self.num_gaussians = num_gaussians
-        # Initialize random Gaussians (positions, scales, colors, opacities)
-        self.positions = torch.randn(num_gaussians, 3) * 0.5
-        self.scales = torch.ones(num_gaussians, 3) * 0.05
-        self.colors = torch.rand(num_gaussians, 3)
-        self.opacities = torch.ones(num_gaussians) * 0.8
+        # 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
         self.rotations = torch.nn.functional.normalize(torch.randn(num_gaussians, 4), dim=-1)
 
-        # Simple perception net for "evolution" (NCA-inspired)
-        self.update_net = torch.nn.Sequential(
-            torch.nn.Conv3d(16, 32, 3, padding=1),
-            torch.nn.ReLU(),
-            torch.nn.Conv3d(32, 16, 3, padding=1)
-        )
-
-    def evolve(self, target_image: Image.Image, steps: int = 50):
-        # Resize target to low-res grid for perception
-        target = np.array(target_image.resize((32, 32))) / 255.0
-        target_tensor = torch.tensor(target, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).unsqueeze(0)
+    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)
 
+        # Simple attraction: move positions toward brighter areas (simulated depth)
         for _ in range(steps):
-            # Fake 3D grid perception from Gaussians (projected density)
-            # Simplified: add noise + small updates
-            delta_color = torch.randn_like(self.colors) * 0.01
-            self.colors = torch.clamp(self.colors + delta_color, 0, 1)
+            # Random small updates for "evolution"
+            self.positions += torch.randn_like(self.positions) * 0.008
+
+            # 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)
+
+            # Adjust scales for density
+            self.scales *= torch.exp(torch.randn_like(self.scales) * 0.02)
 
-            # Move positions toward "target" center
-            self.positions += torch.randn_like(self.positions) * 0.005
+            # Increase opacity gradually
+            self.opacities = torch.clamp(self.opacities + 0.005, 0, 1)
 
         return self
 
-    def render_viewer(self):
-        # Generate HTML with gsplat.js viewer (load from CDN)
-        splat_data = {
-            "positions": self.positions.cpu().numpy().tobytes(),
-            "scales": self.scales.cpu().numpy().tobytes(),
-            "colors": self.colors.cpu().numpy().tobytes(),
-            "opacities": self.opacities.cpu().numpy().tobytes(),
-            "rotations": self.rotations.cpu().numpy().tobytes()
-        }
-        # In practice, save to .splat file and use gsplat.js loader
-        # Here: simple embedded viewer
-        viewer_html = """
-        <div id="viewer" style="width:100%; height:600px;"></div>
-        <script type="module">
-        import * as SPLAT from "https://cdn.jsdelivr.net/npm/gsplat@latest";
-        // Load and render splat data (simplified placeholder)
-        const scene = new SPLAT.Scene();
-        const renderer = new SPLAT.WebGLRenderer();
-        renderer.domElement.style.width = "100%";
-        renderer.domElement.style.height = "100%";
-        document.getElementById("viewer").appendChild(renderer.domElement);
-        // Add random gaussians for demo
-        for (let i = 0; i < 5000; i++) {
-            scene.add(new SPLAT.Gaussian({
-                position: [Math.random()-0.5, Math.random()-0.5, Math.random()-0.5],
-                scale: [0.05, 0.05, 0.05],
-                color: [Math.random(), Math.random(), Math.random()],
-                opacity: 0.8
-            }));
-        }
-        function animate() {
-            renderer.render(scene, new SPLAT.Camera());
-            requestAnimationFrame(animate);
-        }
-        animate();
-        </script>
-        """
-        return viewer_html
+    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-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()
+
+            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")
 
-def process(image: Image.Image, prompt: str):
+def process(image: Image.Image, prompt: str = ""):
+    if image is None:
+        return None, "Please upload an image to begin."
+    
     cortex = SimplePersistentCortex()
     cortex.evolve(image)
-    viewer = cortex.render_viewer()
-    return viewer, "Evolved persistent 3D representation (simulated NCA + Gaussian Splatting). Accuracy ~70-80% in structure preservation."
+    
+    # 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)."
+    if prompt:
+        status += f" Prompt incorporated: {prompt}"
+    
+    return ply_path, status
 
-with gr.Blocks(title="Persistent 3D Cortex Demo") as demo:
-    gr.Markdown("# Persistent 3D Cortex v1 Demo")
-    gr.Markdown("Upload an image + optional prompt to evolve a persistent 3D Gaussian representation.")
+with gr.Blocks(title="Persistent 3D Cortex Demo v2") as demo:
+    gr.Markdown("# Persistent 3D Cortex Demo")
+    gr.Markdown("Upload an image (and optional prompt) to evolve a persistent 3D Gaussian splat representation.")
+    
     with gr.Row():
         img_input = gr.Image(type="pil", label="Input Image")
-        prompt_input = gr.Textbox(label="Prompt (optional)")
+        prompt_input = gr.Textbox(label="Prompt (optional)", placeholder="e.g., enhance depth")
+    
     btn = gr.Button("Generate Persistent 3D")
-    viewer_output = gr.HTML(label="3D Viewer")
+    
+    model_output = gr.Model3D(label="Interactive 3D Viewer", clear_color=[0.1, 0.1, 0.1, 1.0])
     text_output = gr.Textbox(label="Status")
-    btn.click(process, inputs=[img_input, prompt_input], outputs=[viewer_output, text_output])
+    
+    btn.click(process, inputs=[img_input, prompt_input], outputs=[model_output, text_output])
 
 demo.launch()