fix: end-to-end audit — repair source bugs and test suite

.github/workflows/ci.yml

@@ -38,6 +38,10 @@ jobs:
           tests/test_model_transforms.py \
           tests/test_integration.py \
           tests/test_performance.py \
+          tests/test_parallel_process.py \
+          tests/test_colorbar.py \
+          tests/test_entropy_model.py \
+          tests/test_octree_coding.py \
           -v --cov=src --cov-report=xml -m "not gpu and not slow"
 
     - name: Upload coverage

pytest.ini

@@ -3,5 +3,6 @@ markers =
     gpu: marks tests that require a GPU
     e2e: marks end-to-end tests
     integration: marks integration tests
+    slow: marks slow-running tests
 testpaths = tests
 addopts = -ra

src/cli_train.py

@@ -30,8 +30,8 @@ def load_and_preprocess_data(input_dir, batch_size):
     file_paths = glob.glob(os.path.join(input_dir, "*.ply"))
 
     def parse_ply_file(file_path):
-        vertices, _ = read_off(file_path)
-        return vertices
+        mesh_data = read_off(file_path)
+        return mesh_data.vertices
 
     def data_generator():
         for file_path in file_paths:
@@ -82,7 +82,14 @@ def main():
     if args.tune:
         tune_hyperparameters(args.input_dir, args.output_dir, num_epochs=args.num_epochs)
     else:
-        model = create_model(hp=None)
+        # Build a default model without hyperparameter tuning
+        model = tf.keras.Sequential([
+            tf.keras.layers.InputLayer(input_shape=(2048, 3)),
+            tf.keras.layers.Dense(256, activation='relu'),
+            tf.keras.layers.Dense(128, activation='relu'),
+            tf.keras.layers.Dense(3, activation='sigmoid')
+        ])
+        model.compile(optimizer='adam', loss='mean_squared_error')
         dataset = load_and_preprocess_data(args.input_dir, args.batch_size)
         model.fit(dataset, epochs=args.num_epochs)
         model.save(os.path.join(args.output_dir, 'trained_model'))

src/compress_octree.py

@@ -75,6 +75,9 @@ def _create_normal_grid(self, indices: np.ndarray, normals: np.ndarray) -> np.nd
 
     def compress(self, point_cloud: np.ndarray, normals: Optional[np.ndarray] = None, validate: bool = True) -> Tuple[np.ndarray, Dict[str, Any]]:
         """Compress point cloud with optional normals and validation."""
+        point_cloud = np.asarray(point_cloud)
+        if normals is not None:
+            normals = np.asarray(normals)
         if len(point_cloud) == 0:
             raise ValueError("Empty point cloud provided")
 
@@ -117,6 +120,7 @@ def decompress(self, grid: np.ndarray, metadata: Dict[str, Any], *, return_norma
 
     def partition_octree(self, point_cloud: np.ndarray, max_points_per_block: int = 1000, min_block_size: float = 0.1) -> List[Tuple[np.ndarray, Dict[str, Any]]]:
         """Partition point cloud into octree blocks."""
+        point_cloud = np.asarray(point_cloud)
         blocks = []
         min_bounds = np.min(point_cloud, axis=0)
         max_bounds = np.max(point_cloud, axis=0)

src/data_loader.py

@@ -1,4 +1,5 @@
 import tensorflow as tf
+import numpy as np
 import glob
 import os
 from pathlib import Path
@@ -16,20 +17,23 @@ def __init__(self, config: Dict[str, Any]):
             min_points=config.get('min_points', 100)
         )
 
-    @tf.function
     def process_point_cloud(self, file_path: str) -> tf.Tensor:
         """Process a single point cloud file."""
         # Read point cloud
-        vertices, _ = read_off(file_path.numpy().decode())
-        points = tf.convert_to_tensor(vertices, dtype=tf.float32)
-        
+        mesh_data = read_off(file_path.numpy().decode())
+        points = tf.convert_to_tensor(mesh_data.vertices, dtype=tf.float32)
+
         # Normalize points to unit cube
         points = self._normalize_points(points)
-
+
+        # Apply augmentation before voxelization
+        if self.config.get('augment', True):
+            points = self._augment(points)
+
         # Voxelize points
         resolution = self.config.get('resolution', 64)
         voxelized = self._voxelize_points(points, resolution)
-        
+
         return voxelized
 
     @tf.function
@@ -75,13 +79,6 @@ def load_training_data(self) -> tf.data.Dataset:
             num_parallel_calls=tf.data.AUTOTUNE
         )
 
-        # Apply training augmentations
-        if self.config.get('augment', True):
-            dataset = dataset.map(
-                self._augment,
-                num_parallel_calls=tf.data.AUTOTUNE
-            )
-
         # Batch and prefetch
         dataset = dataset.shuffle(1000)
         dataset = dataset.batch(self.config['training']['batch_size'])
@@ -123,9 +120,12 @@ def _augment(self, points: tf.Tensor) -> tf.Tensor:
         ])
         points = tf.matmul(points, rotation)
 
-        # Random jittering
-        if tf.random.uniform([]) < 0.5:
-            jitter = tf.random.normal(tf.shape(points), mean=0.0, stddev=0.01)
-            points = points + jitter
-
+        # Random jittering (use tf.cond for @tf.function compatibility)
+        jitter = tf.random.normal(tf.shape(points), mean=0.0, stddev=0.01)
+        points = tf.cond(
+            tf.random.uniform([]) < 0.5,
+            lambda: points + jitter,
+            lambda: points
+        )
+
         return points

src/ds_pc_octree_blocks.py

@@ -28,7 +28,6 @@ def parse_line(line):
         )
         return points
 
-    @tf.function
     def partition_point_cloud(self, points: tf.Tensor) -> List[tf.Tensor]:
         """Partition point cloud into blocks using TF operations."""
         # Compute bounds

src/ds_select_largest.py

@@ -14,13 +14,14 @@ def count_points_in_block(file_path):
     """
     with open(file_path, "r") as f:
         header = True
+        count = 0
         for line in f:
             if header:
                 if line.startswith("end_header"):
                     header = False
                 continue
-        # Count remaining lines
-        return sum(1 for _ in f)
+            count += 1
+        return count
 
 def prioritize_blocks(input_dir, output_dir, num_blocks, criteria="points"):
     """

src/entropy_model.py

@@ -113,7 +113,6 @@ def quantize_scale(self, scale: tf.Tensor) -> tf.Tensor:
         quantized_flat = tf.gather(self.scale_table, indices)
         return tf.reshape(quantized_flat, original_shape)
 
-    @tf.function
     def compress(self, inputs: tf.Tensor) -> tf.Tensor:
         scale = self.quantize_scale(self.scale)
         centered = inputs - self.mean
@@ -128,7 +127,6 @@ def compress(self, inputs: tf.Tensor) -> tf.Tensor:
 
         return quantized
 
-    @tf.function
     def decompress(self, inputs: tf.Tensor) -> tf.Tensor:
         scale = self.quantize_scale(self.scale)
         denormalized = inputs * scale
@@ -142,7 +140,6 @@ def decompress(self, inputs: tf.Tensor) -> tf.Tensor:
 
         return decompressed
 
-    @tf.function
     def call(self, inputs: tf.Tensor, training: Optional[bool] = None) -> tf.Tensor:
         self._debug_tensors['inputs'] = inputs
         compressed = self.compress(inputs)
@@ -209,7 +206,6 @@ def _add_noise(self, inputs: tf.Tensor, training: bool) -> tf.Tensor:
             return inputs + noise
         return tf.round(inputs)
 
-    @tf.function
     def compress(self, inputs: tf.Tensor, scale: tf.Tensor, mean: tf.Tensor) -> tf.Tensor:
         """
         Compress inputs using provided scale and mean.
@@ -238,7 +234,6 @@ def compress(self, inputs: tf.Tensor, scale: tf.Tensor, mean: tf.Tensor) -> tf.T
 
         return quantized
 
-    @tf.function
     def decompress(self, inputs: tf.Tensor, scale: tf.Tensor, mean: tf.Tensor) -> tf.Tensor:
         """
         Decompress inputs using provided scale and mean.

src/evaluation_pipeline.py

@@ -3,7 +3,7 @@
 import logging
 from pathlib import Path
 from typing import Dict, Any, List
-from dataclasses import dataclass
+from dataclasses import dataclass, asdict
 
 from data_loader import DataLoader
 from model_transforms import DeepCompressModel, TransformConfig
@@ -55,61 +55,62 @@ def _load_model(self) -> DeepCompressModel:
 
         return model
 
-    @tf.function
-    def _evaluate_single(self, 
+    def _evaluate_single(self,
                         point_cloud: tf.Tensor) -> Dict[str, tf.Tensor]:
         """Evaluate model on single point cloud."""
-        # Forward pass
-        compressed, metrics = self.model.compress(point_cloud)
-        decompressed = self.model.decompress(compressed)
-
+        # Forward pass through model
+        x_hat, y, y_hat, z = self.model(point_cloud, training=False)
+
         # Compute metrics
         results = {}
-        results['psnr'] = self.metrics.compute_psnr(point_cloud, decompressed)
-        results['chamfer'] = self.metrics.compute_chamfer(point_cloud, decompressed)
-
-        # Add compression metrics
-        results.update(metrics)
-
+        results['psnr'] = self.metrics.compute_psnr(point_cloud, x_hat)
+        results['chamfer'] = self.metrics.compute_chamfer(point_cloud, x_hat)
+
         return results
 
     def evaluate(self) -> Dict[str, List[EvaluationResult]]:
         """Run evaluation on test dataset."""
         results = {}
         dataset = self.data_loader.load_evaluation_data()
 
-        for point_cloud in dataset:
-            # Get filename from dataset
-            filename = point_cloud.filename.numpy().decode('utf-8')
+        for i, point_cloud in enumerate(dataset):
+            filename = f"point_cloud_{i}"
             self.logger.info(f"Evaluating {filename}")
-            
+
             try:
                 # Time compression and decompression
                 start_time = tf.timestamp()
                 metrics = self._evaluate_single(point_cloud)
                 end_time = tf.timestamp()
-                
+
                 # Create result object
                 result = EvaluationResult(
                     psnr=float(metrics['psnr']),
                     chamfer_distance=float(metrics['chamfer']),
                     bd_rate=float(metrics.get('bd_rate', 0.0)),
-                    file_size=int(metrics['compressed_size']),
+                    file_size=int(metrics.get('compressed_size', 0)),
                     compression_time=float(end_time - start_time),
                     decompression_time=float(metrics.get('decompress_time', 0.0))
                 )
-                
+
                 results[filename] = result
-                
+
             except Exception as e:
                 self.logger.error(f"Error processing {filename}: {str(e)}")
                 continue
 
         return results
 
-    def generate_report(self, results: Dict[str, List[EvaluationResult]]):
+    def generate_report(self, results: Dict[str, EvaluationResult]):
         """Generate evaluation report."""
-        reporter = ExperimentReporter(results)
+        # Convert EvaluationResult objects to flat dicts for ExperimentReporter
+        flat_results = {}
+        for name, result in results.items():
+            if isinstance(result, EvaluationResult):
+                flat_results[name] = asdict(result)
+            else:
+                flat_results[name] = result
+        reporter = ExperimentReporter(flat_results)
 
         # Generate and save report
         output_dir = Path(self.config['evaluation']['output_dir'])

src/map_color.py

@@ -16,6 +16,9 @@ def load_point_cloud(file_path: str) -> Optional[np.ndarray]:
     try:
         with open(file_path, 'r') as file:
             lines = file.readlines()
+            if "end_header\n" not in lines:
+                print(f"Error: no end_header found in {file_path}")
+                return None
             start = lines.index("end_header\n") + 1
             points = [list(map(float, line.split()[:3])) for line in lines[start:]]
             return np.array(points, dtype=np.float32)
@@ -36,6 +39,9 @@ def load_colors(file_path: str) -> Optional[np.ndarray]:
     try:
         with open(file_path, 'r') as file:
             lines = file.readlines()
+            if "end_header\n" not in lines:
+                print(f"Error: no end_header found in {file_path}")
+                return None
             start = lines.index("end_header\n") + 1
             colors = [list(map(float, line.split()[3:6])) for line in lines[start:] if len(line.split()) > 3]
             return np.array(colors, dtype=np.float32) / 255.0

src/model_transforms.py

@@ -2,7 +2,7 @@
 from typing import Tuple
 from dataclasses import dataclass
 
-from constants import LOG_2_RECIPROCAL
+from constants import LOG_2_RECIPROCAL, EPSILON
 
 
 @dataclass
@@ -43,7 +43,7 @@ def call(self, x):
         # Use axis 4 (channel dimension) for 5D tensors (batch, D, H, W, C)
         norm = tf.tensordot(norm, self.gamma, [[4], [0]])
         norm = tf.nn.bias_add(norm, self.beta)
-        return x / norm
+        return x / tf.maximum(norm, EPSILON)
 
     def get_config(self):
         config = super().get_config()
@@ -208,6 +208,11 @@ def __init__(self, config: TransformConfig, **kwargs):
         self.analysis = AnalysisTransform(config)
         self.synthesis = SynthesisTransform(config)
 
+        # Final projection: map from synthesis channels back to 1-channel occupancy
+        self.output_projection = tf.keras.layers.Conv3D(
+            filters=1, kernel_size=(1, 1, 1), activation='sigmoid', padding='same'
+        )
+
         # Hyperprior
         self.hyper_analysis = AnalysisTransform(TransformConfig(
             filters=config.filters // 2,
@@ -232,7 +237,7 @@ def call(self, inputs, training=None):
 
         # Synthesis
         y_hat = self.hyper_synthesis(z)
-        x_hat = self.synthesis(y)
+        x_hat = self.output_projection(self.synthesis(y))
 
         return x_hat, y, y_hat, z
 
@@ -288,6 +293,11 @@ def __init__(self,
         self.analysis = AnalysisTransform(config)
         self.synthesis = SynthesisTransform(config)
 
+        # Final projection: map from synthesis channels back to 1-channel occupancy
+        self.output_projection = tf.keras.layers.Conv3D(
+            filters=1, kernel_size=(1, 1, 1), activation='sigmoid', padding='same'
+        )
+
         # Hyperprior transforms
         self.hyper_analysis = AnalysisTransform(TransformConfig(
             filters=config.filters // 2,
@@ -400,7 +410,7 @@ def call(self, inputs, training=None):
             )
 
         # Synthesis
-        x_hat = self.synthesis(y_hat)
+        x_hat = self.output_projection(self.synthesis(y_hat))
 
         # Rate information
         rate_info = {
@@ -474,7 +484,7 @@ def decompress(self, compressed_data):
             y_hat = y_compressed
 
         # Synthesis
-        x_hat = self.synthesis(y_hat)
+        x_hat = self.output_projection(self.synthesis(y_hat))
 
         return x_hat