Update app.py

app.py

@@ -1,23 +1,23 @@
 import laspy
 import numpy as np
 import cv2
-from math import floor, pi
+from math import pi
 from multiprocessing import Process, shared_memory, cpu_count
 import gradio as gr
-import time
 
-# ==== Constants ====
-TILE_DIM = 512
+# === Constants ===
 PRECISION = 16
 FACTOR = 2 ** PRECISION
-NUM_WORKERS = cpu_count()
+NUM_WORKERS = min(cpu_count(), 4)
 SHARED_MEMORY_NAME = 'shared_canvas'
 
-# ==== Parameters ====
-IMAGE_HEIGHT = 4096 * 2
+# === Parameters ===
+IMAGE_HEIGHT = 4096 * 2  # Default image height
 POINT_SIZE = 3
 SCALE_COLORS = False
 FULL_RES = False
+MAX_RADIUS = 30
+
 
 def create_shared_memory_array(data, name, dtype):
     shm = shared_memory.SharedMemory(create=True, size=data.nbytes, name=name)
@@ -25,85 +25,110 @@ def create_shared_memory_array(data, name, dtype):
     shared_array[:] = data
     return shm
 
+
 def release_shared_memory(name):
     shm = shared_memory.SharedMemory(name=name)
     shm.close()
     shm.unlink()
 
+
 def draw_points(start, length, canvas_shape, theta, phi, radius, color):
     shm = shared_memory.SharedMemory(name=SHARED_MEMORY_NAME)
     canvas = np.ndarray(canvas_shape, dtype=np.uint16, buffer=shm.buf)
 
     for i in range(start, start + length):
-        cv2.circle(
-            canvas,
-            (theta[i], phi[i]),
-            radius[i],
-            color=(int(color[i][0]), int(color[i][1]), int(color[i][2])),
-            thickness=cv2.FILLED,
-            shift=PRECISION
-        )
+        if 0 <= theta[i] < canvas_shape[1] * FACTOR and 0 <= phi[i] < canvas_shape[0] * FACTOR:
+            cv2.circle(
+                canvas,
+                (int(theta[i]), int(phi[i])),
+                int(radius[i]),
+                color=(int(color[i][0]), int(color[i][1]), int(color[i][2])),
+                thickness=cv2.FILLED,
+                shift=PRECISION
+            )
+
 
 def project_point_cloud(input_file):
     try:
         point_cloud = laspy.read(input_file.name)
-    except FileNotFoundError:
-        return "Input file not found."
+    except Exception as e:
+        return f"Failed to read LAS file: {e}"
 
     height, width = IMAGE_HEIGHT, IMAGE_HEIGHT * 2
     output_shape = (height, width, 3)
+
     total_points = point_cloud.header.point_count
+    print(f"[INFO] Loaded {total_points} points.")
+
+    canvas = np.zeros(output_shape, dtype=np.uint16)
 
-    background = np.zeros(output_shape, dtype=np.uint16)
+    # === Coordinate setup ===
+    xyz = np.vstack((point_cloud.x, point_cloud.y, point_cloud.z)).T
+    distances = np.linalg.norm(xyz, axis=1)
+    distances[distances == 0] = 1e-6  # Avoid divide by zero
 
-    xyz = np.array([point_cloud.x, point_cloud.y, point_cloud.z])
-    distances = np.linalg.norm(xyz, axis=0)
-    distances_norm = (distances - distances.min()) / (distances.max() - distances.min())
+    # === Normalize distances for radius scaling ===
+    dist_norm = (distances - distances.min()) / (distances.max() - distances.min())
 
-    theta = ((width * FACTOR) -
-             (((np.arctan2(point_cloud.y, point_cloud.x) + pi) / (2 * pi)) * width * FACTOR)).astype(np.uint64)
-    phi = ((np.arccos(point_cloud.z / distances) / pi) * height * FACTOR).astype(np.uint64)
-    radii = (FACTOR * POINT_SIZE * (1 - distances_norm) + 1).astype(np.uint64)
+    # === Compute angles ===
+    theta = ((np.arctan2(point_cloud.y, point_cloud.x) + pi) / (2 * pi)) * width * FACTOR
+    theta = np.mod(theta, width * FACTOR).astype(np.uint64)
 
+    z_norm = np.clip(point_cloud.z / distances, -1.0, 1.0)
+    phi = ((np.arccos(z_norm) / pi) * height * FACTOR).astype(np.uint64)
+
+    # === Radius based on depth (closer → larger radius) ===
+    radii = (FACTOR * POINT_SIZE * (1 - dist_norm) + 1).astype(np.uint64)
+    radii = np.clip(radii, 1, MAX_RADIUS)
+
+    # === Colors ===
     colors = np.stack([point_cloud.blue, point_cloud.green, point_cloud.red], axis=-1).astype(np.uint16)
     if SCALE_COLORS:
         colors *= 256
+    colors = np.clip(colors, 0, 65535)
 
-    shm_canvas = create_shared_memory_array(background, SHARED_MEMORY_NAME, np.uint16)
+    # === Shared Memory Canvas ===
+    shm = create_shared_memory_array(canvas, SHARED_MEMORY_NAME, np.uint16)
 
+    # === Start workers ===
     processes = []
-    num_threads = min(NUM_WORKERS, 4)
-    for i in range(num_threads):
-        start = floor(i * total_points / num_threads)
-        end = floor((i + 1) * total_points / num_threads)
-        p = Process(target=draw_points, args=(start, end - start, output_shape, theta, phi, radii, colors))
+    for i in range(NUM_WORKERS):
+        start = int(i * total_points / NUM_WORKERS)
+        end = int((i + 1) * total_points / NUM_WORKERS)
+        p = Process(target=draw_points, args=(
+            start, end - start, output_shape, theta, phi, radii, colors
+        ))
         p.start()
         processes.append(p)
 
     for p in processes:
         p.join()
 
-    final_image = np.ndarray(output_shape, dtype=np.uint16, buffer=shm_canvas.buf)
+    # === Finalize output ===
+    final_image = np.ndarray(output_shape, dtype=np.uint16, buffer=shm.buf)
     if not FULL_RES:
         final_image = (final_image / 256).astype(np.uint8)
 
     output_file = "output_image.jpg"
     cv2.imwrite(output_file, final_image)
-
     release_shared_memory(SHARED_MEMORY_NAME)
+    print(f"[INFO] Saved output to {output_file}")
+
     return output_file
 
-# Gradio interface
+
+# === Gradio Interface ===
 def main(input_file):
     return project_point_cloud(input_file)
 
+
 iface = gr.Interface(
     fn=main,
     inputs=gr.File(label="Upload LAS File"),
-    outputs=gr.Image(type="filepath", label="Output Image"),
-    title="Point Cloud Projection",
-    description="Upload a LAS file to project the point cloud into an image."
+    outputs=gr.Image(type="filepath", label="Projected Image"),
+    title="Equirectangular Point Cloud Projection",
+    description="Upload a LAS point cloud file and project it into a 360° equirectangular image."
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()