Update app.py

app.py

@@ -30,32 +30,37 @@ class PersistentCortex:
 
     def evolve_from_image(self, image: Image.Image, steps=800):
         img_array = np.array(image.resize((256, 256))) / 255.0
-        target = torch.tensor(img_array, dtype=torch.float32).permute(2, 0, 1)  # (3, 256, 256)
+        target = torch.tensor(img_array, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0)  # (1, 3, 256, 256)
 
-        # Create sampling grid: scatter points across [-1,1] x [-1,1]
-        proj = torch.rand(self.num, 2) * 2 - 1  # (N, 2) uniform in view
+        # Random projection points in [-1,1] for sampling
+        proj = torch.rand(self.num, 2) * 2 - 1  # (N, 2)
 
         for _ in range(steps):
+            # Normalize grid to [-1,1]
             grid = proj.unsqueeze(0).unsqueeze(0)  # (1, 1, N, 2)
 
             sampled = torch.nn.functional.grid_sample(
-                target.unsqueeze(0),
+                target,
                 grid,
                 mode='bilinear',
                 padding_mode='border',
                 align_corners=True
-            ).squeeze(0).squeeze(0).t()  # (N, 3)
+            ).squeeze(1).squeeze(0)  # Squeeze H (1) and batch if needed, to (3, N)
+
+            # Transpose to (N, 3)
+            sampled = sampled.t()
 
             brightness = sampled.mean(dim=1)
-            attraction_strength = (brightness - brightness.mean()) * 0.02
-            proj += attraction_strength.unsqueeze(1) * proj.normalized()
+            attraction = (brightness - brightness.mean()) * 0.015
+            proj += attraction.unsqueeze(1) * (proj / (proj.norm(dim=1, keepdim=True) + 1e-6))
 
-            # Color adaptation
-            avg_color = target.mean(dim=[1,2])
+            # Color lerp to image average
+            avg_color = target.mean(dim=[2,3]).squeeze()
             self.colors = torch.lerp(self.colors, avg_color.repeat(self.num, 1), 0.02)
 
-        # Map projected points back to 3D sphere surface (simple radial projection)
-        self.positions[:, :2] = proj * (self.positions[:, :2].norm(dim=1).unsqueeze(1) + 0.1)
+        # Apply evolved projection to positions (scale to sphere radius)
+        radii = self.positions.norm(dim=1, keepdim=True)
+        self.positions[:, :2] = proj * radii
 
         return self
 
@@ -67,6 +72,7 @@ class PersistentCortex:
             text_emb = model.encode_text(text_tokens).float()
             text_emb = text_emb / text_emb.norm(dim=-1, keepdim=True)
 
+        # Projection to color shift
         proj = torch.nn.Linear(512, 3, bias=False)
         torch.nn.init.normal_(proj.weight, std=0.2)
         color_shift = proj(text_emb)
@@ -75,16 +81,20 @@ class PersistentCortex:
         return self
 
     def export_ply(self, path="output.ply"):
-        zeros = np.zeros((self.num, 3), dtype=np.float32)
-        log_scales = np.log(self.scales.cpu().numpy())
+        pos = self.positions.cpu().numpy()
+        col = self.colors.cpu().numpy()
+        opa = self.opacities.cpu().numpy()
+        sca = np.log(self.scales.cpu().numpy())
+        rot = self.rotations.cpu().numpy()
+        nor = np.zeros_like(pos)  # Normals placeholder
 
         vertex_data = np.core.records.fromarrays([
-            self.positions.cpu().numpy(),
-            zeros,
-            self.colors.cpu().numpy(),
-            self.opacities.cpu().numpy()[:, np.newaxis],
-            log_scales,
-            self.rotations.cpu().numpy()
+            pos[:,0], pos[:,1], pos[:,2],
+            nor[:,0], nor[:,1], nor[:,2],
+            col[:,0], col[:,1], col[:,2],
+            opa,
+            sca[:,0], sca[:,1], sca[:,2],
+            rot[:,0], rot[:,1], rot[:,2], rot[:,3]
         ], names='x,y,z,nx,ny,nz,f_dc_0,f_dc_1,f_dc_2,opacity,scale_0,scale_1,scale_2,rot_0,rot_1,rot_2,rot_3')
 
         el = plyfile.PlyElement.describe(vertex_data, 'vertex')
@@ -93,7 +103,7 @@ class PersistentCortex:
 
 def process(image: Image.Image, prompt: str = ""):
     if image is None:
-        raise gr.Error("Please upload an image.")
+        raise gr.Error("Please upload an image to proceed.")
     
     cortex = PersistentCortex(num_gaussians=8000)
     cortex.evolve_from_image(image, steps=800)
@@ -121,6 +131,9 @@ def process(image: Image.Image, prompt: str = ""):
     controls.autoRotate = false;
     controls.enableDamping = true;
     controls.dampingFactor = 0.05;
+    controls.rotateSpeed = 1.0;
+    controls.zoomSpeed = 1.2;
+    controls.panSpeed = 0.8;
 
     await SPLAT.Loader.LoadAsync("/files/output.ply", scene);
 
@@ -133,7 +146,7 @@ def process(image: Image.Image, prompt: str = ""):
     </script>
     """
 
-    status = f"Persistent 3D representation evolved from image"
+    status = "Persistent 3D representation evolved from image"
     if prompt:
         status += f" and conditioned on prompt: '{prompt}'"
     status += "."
@@ -145,14 +158,18 @@ with gr.Blocks(title="Persistent 3D Cortex Demo") as demo:
     gr.Markdown("Upload an image and optionally add a text prompt to generate an evolving 3D Gaussian splat representation.")
     
     with gr.Row():
-        img_input = gr.Image(type="pil", label="Input Image (required)")
-        prompt_input = gr.Textbox(label="Prompt (e.g., 'glowing blue crystal')", placeholder="Optional descriptive text")
+        img_input = gr.Image(type="pil", label="Input Image")
+        prompt_input = gr.Textbox(label="Prompt (e.g., 'shiny red apple', 'futuristic city')", placeholder="Optional text prompt")
     
-    generate_btn = gr.Button("Generate Persistent 3D", variant="primary")
+    generate_btn = gr.Button("Generate & Evolve 3D", variant="primary")
     
-    viewer_output = gr.HTML()
+    viewer_output = gr.HTML(label="Interactive 3D Viewer")
     status_output = gr.Textbox(label="Status")
 
-    generate_btn.click(process, inputs=[img_input, prompt_input], outputs=[viewer_output, status_output])
+    generate_btn.click(
+        fn=process,
+        inputs=[img_input, prompt_input],
+        outputs=[viewer_output, status_output]
+    )
 
 demo.launch()