diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e9e198d56..d5a94fdb3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -7,8 +7,20 @@ on: branches: [main] jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.10' + - run: pip install ruff + - name: Ruff check + run: ruff check src/ tests/ + test: runs-on: ubuntu-latest + needs: lint strategy: matrix: python-version: ['3.8', '3.9', '3.10'] @@ -29,7 +41,6 @@ jobs: - name: Run tests run: | - # Run entropy modeling tests and performance tests (legacy tests have pre-existing issues) pytest \ tests/test_entropy_parameters.py \ tests/test_context_model.py \ @@ -48,36 +59,3 @@ jobs: uses: codecov/codecov-action@v4 with: file: ./coverage.xml - - lint: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 - with: - python-version: '3.10' - - run: pip install flake8 - # Lint only the new entropy modeling and optimization files (legacy files have pre-existing issues) - - name: Lint new source files - run: | - flake8 \ - src/entropy_parameters.py \ - src/entropy_model.py \ - src/context_model.py \ - src/channel_context.py \ - src/attention_context.py \ - src/model_transforms.py \ - src/constants.py \ - src/precision_config.py \ - src/benchmarks.py \ - --max-line-length=120 - - name: Lint new test files - run: | - flake8 \ - tests/test_entropy_parameters.py \ - tests/test_context_model.py \ - tests/test_channel_context.py \ - tests/test_attention_context.py \ - tests/test_performance.py \ - --max-line-length=120 \ - --ignore=E402,W503 # E402: imports after sys.path, W503: PEP8 updated to prefer breaks before operators diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 000000000..2f7cab7aa --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,388 @@ +# CLAUDE.md — DeepCompress + +> This file is for AI coding agents working on this codebase. It defines how to write, test, and modify code here without introducing the mistakes that LLMs typically make. + +## What This Project Is + +DeepCompress is a TensorFlow 2 library for learned point cloud compression. It compresses 3D voxelized occupancy grids using neural analysis/synthesis transforms with configurable entropy models. **It is not a PyTorch project.** + +Python 3.8+ · TensorFlow ≥ 2.11 · tensorflow-probability ~0.19 · MIT License + +## Before You Write Any Code + +1. **Read the existing code first.** Before modifying a file, read it entirely. Before adding a function, check if a similar one exists. Before creating a test utility, check `tests/test_utils.py` — it likely already has what you need. + +2. **Run the tests.** Before and after every change: `pytest tests/ -v`. Do not submit changes that break existing tests. If a test was already skipped/broken before your change, leave it that way — do not fix unrelated tests as part of your change. + +3. **Make minimal changes.** Fix what was asked. Do not refactor adjacent code, rename variables for "clarity," add type hints to unrelated functions, or restructure files. Every unnecessary change is a potential regression. + +--- + +## Agent Workflow Rules + +### Change Verification Loop + +Every code change must follow this cycle: + +``` +1. Read the file(s) you plan to modify +2. Make the smallest change that addresses the task +3. Run pytest — at minimum the test file for the module you changed +4. If tests fail, fix your change (not the tests) unless the tests are genuinely wrong +5. Run full pytest to check for regressions +``` + +### Do Not + +- Edit multiple source files at once without testing between edits +- "Improve" code you weren't asked to touch +- Add logging, comments, or docstrings unless specifically requested +- Create new abstractions (base classes, mixins, utility modules) without being asked +- Remove or weaken existing error handling, assertions, or numerical guards +- Replace working code with "cleaner" alternatives that change behavior + +--- + +## Framework Rules: This Is TensorFlow + +This is the single most common LLM mistake on this codebase. Every line below matters. + +- **All layers subclass `tf.keras.layers.Layer`**. Models subclass `tf.keras.Model`. There is no `nn.Module`, no `.forward()`, no `torch.Tensor`. +- **Trainable state goes in `build()`.** Use `self.add_weight()` for learned parameters. Non-trainable configuration state (like `scale_table` in `PatchedGaussianConditional`) may be set in `__init__()` as a `tf.Variable(..., trainable=False)` — this is an established pattern in the codebase. +- **Do not create `tf.Variable` or `tf.constant` inside `call()`.** This leaks memory under `@tf.function` tracing. All persistent state must be initialized in `__init__()` or `build()`. +- **Use `get_config()`** for serialization on Keras layers that participate in model graphs. Standalone utility classes (like `OctreeCoder`) do not need it. +- **`@tf.function` compatibility.** Code in `call()` must work in graph mode. No Python-level `if tensor_value > 0:` — use `tf.cond`. No creating new ops that depend on Python-side tensor evaluation. +- **Never import torch, torch.nn, or any PyTorch module.** Do not suggest PyTorch alternatives. + +### The `training` flag + +The main model `call()` methods (`DeepCompressModel.call()`, `DeepCompressModelV2.call()`) accept and use `training=None`. This is required because quantization behavior switches between noise injection (training) and hard rounding (inference): + +```python +# This pattern appears in DeepCompressModel.call(), ConditionalGaussian._add_noise(), etc. +if training: + y = y + tf.random.uniform(tf.shape(y), -0.5, 0.5) # Gradient-friendly noise +else: + y = tf.round(y) # Hard quantization for inference +``` + +**Rules:** +- Model-level `call()` methods and any layer that branches on training mode **must** accept `training=None` and pass it through to sub-layers that need it. +- Leaf layers that do not use `training` internally (e.g., `CENICGDN`, `SpatialSeparableConv`, `MaskedConv3D`, `SliceTransform`) currently omit it from their signatures. This is the established convention — do not add `training` to these unless they gain training-dependent behavior. +- **Never remove the training conditional** from methods that have it. Never replace noise injection with unconditional `tf.round()`. +- When adding new layers: include `training=None` if the layer has any training-dependent behavior. Omit it for pure computation layers. + +## Tensor Shape Convention + +All model tensors are 5D: `(batch, depth, height, width, channels)` — channels-last. + +- Convolutions are `Conv3D`, never `Conv2D`. Kernels are 3-tuples: `(3, 3, 3)`. +- Channel axis is axis 4 (see `CENICGDN.call()` which does `tf.tensordot(norm, self.gamma, [[4], [0]])`). +- Input voxel grids have 1 channel: shape `(B, D, H, W, 1)`. +- Do not flatten spatial dimensions to use 2D ops. The 3D structure is load-bearing. + +## Constants and Numerical Stability + +**Use pre-computed constants from `src/constants.py`:** +```python +from constants import LOG_2_RECIPROCAL +bits = -log_likelihood * LOG_2_RECIPROCAL # CORRECT +bits = -log_likelihood / tf.math.log(2.0) # WRONG: creates new op node every call +``` + +**Do not remove numerical guards.** The codebase uses: +- `tf.maximum(scale, self.scale_min)` — prevents division by zero in Gaussian likelihood +- `EPSILON = 1e-9` — prevents log(0) +- `tf.clip_by_value` on scales — prevents NaN in entropy computation +- `tf.abs(scale)` before quantization — ensures positive scale + +These look like they could be "simplified away." They cannot. Removing them causes NaN gradients during training. + +## Entropy Model Contracts + +The six entropy models (`gaussian`, `hyperprior`, `channel`, `context`, `attention`, `hybrid`) are not interchangeable at runtime. Each is selected as a string at model construction and creates different submodules. Key contracts: + +- `DeepCompressModel.call()` returns 4 values: `(x_hat, y, y_hat, z)` +- `DeepCompressModelV2.call()` returns 5 values: `(x_hat, y, y_hat, z, rate_info)` +- `rate_info` is a dict with keys: `likelihood`, `total_bits`, `bpp` +- Don't mix V1 and V2 return signatures +- Entropy model submodules use deferred imports inside `_create_entropy_model()` to avoid circular dependencies. Keep this pattern. + +## Binary Search Scale Quantization + +`PatchedGaussianConditional.quantize_scale()` uses `tf.searchsorted` on pre-computed midpoints. This is an intentional optimization (O(log T) vs O(T)). + +- Do not replace with linear scan or full-table broadcasting +- The scale table and midpoints are created once in `__init__`/`build()`, not per call +- `_precompute_midpoints()` must be called after any scale table change + +## Masked Convolutions Are Causal + +`context_model.py` has `MaskedConv3D` with type A (excludes center) and type B (includes center) masks. These enforce autoregressive ordering in raster-scan order over (D, H, W). + +- Masks are created via vectorized NumPy operations in `build()`, stored as weights +- Do not convert to Python loops — the vectorized version is 10-100x faster +- Do not modify the raster-scan ordering — it will silently produce incorrect likelihoods + +--- + +## Testing Standards + +### Use the Existing Test Infrastructure + +The test suite uses `tf.test.TestCase` mixed with `pytest` fixtures. This is intentional and established — do not convert to pure pytest. + +**Existing utilities in `tests/test_utils.py` — use these instead of creating your own:** +- `create_mock_voxel_grid(resolution, batch_size)` — binary occupancy grid +- `create_mock_point_cloud(num_points)` — random 3D points +- `create_test_dataset(batch_size, resolution, num_batches)` — tf.data.Dataset +- `create_test_config(tmp_path)` — complete YAML config dict +- `setup_test_environment(tmp_path)` — full test env with files, configs, dirs +- `MockCallback` — for testing training loops + +**Existing fixtures in `tests/conftest.py`:** +- `sample_point_cloud` (session-scoped) — 8-point cube +- `create_ply_file`, `create_off_file` — file factories +- `tf_config` (session-scoped) — GPU memory growth + seed setting + +### How to Write Tests + +**Test structure:** Tests are organized by module. `test_foo.py` tests `src/foo.py`. Follow the existing pattern: + +```python +import tensorflow as tf +import pytest +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) + +from test_utils import create_mock_voxel_grid +from module_under_test import ClassUnderTest + +class TestClassUnderTest(tf.test.TestCase): + @pytest.fixture(autouse=True) + def setup(self): + # Use small tensors: resolution 16, batch_size 1, filters 32 + self.config = TransformConfig(filters=32, ...) + self.resolution = 16 + self.batch_size = 1 + + def test_specific_behavior(self): + # Arrange + input_tensor = create_mock_voxel_grid(self.resolution, self.batch_size) + model = ClassUnderTest(self.config) + + # Act + output = model(input_tensor, training=False) + + # Assert — use tf.test.TestCase assertions + self.assertEqual(output.shape[:-1], input_tensor.shape[:-1]) + self.assertGreater(rate_info['total_bits'], 0) +``` + +**Test sizing:** Use small tensors. Resolution 16, batch size 1, filters 32 is the standard for unit tests. Tests using resolution 64 or larger should be marked `@pytest.mark.slow`. + +**Random seeds:** Set `tf.random.set_seed(42)` at the start of any test that creates random tensors and makes value-dependent assertions. The session-level seed in `conftest.py` provides baseline reproducibility, but per-test seeds are more robust for value-dependent checks. + +**Pytest markers** (defined in `pytest.ini`): +- `@pytest.mark.slow` — performance/timing tests +- `@pytest.mark.gpu` — requires GPU +- `@pytest.mark.e2e` — end-to-end pipeline tests +- `@pytest.mark.integration` — integration tests + +### What Makes a Good Assertion + +**Good assertions test meaningful properties:** +```python +# Shape is correct and specific +self.assertEqual(output.shape, (1, 16, 16, 16, 128)) + +# Values are in expected range +self.assertGreater(rate_info['total_bits'], 0) + +# Roundtrip consistency +self.assertEqual(decompressed.shape, x_hat.shape) + +# Gradients actually flow and are non-zero +non_none_grads = [g for g in gradients if g is not None] +self.assertNotEmpty(non_none_grads) +self.assertGreater(tf.reduce_sum(tf.abs(gradients[0])), 0) + +# Structural correctness (mask causality) +assert np.all(mask[2, 1, 1, :, :] == 0), "Future positions should be masked" + +# Numerical correctness with appropriate tolerances +np.testing.assert_allclose(bits_original.numpy(), bits_optimized.numpy(), rtol=1e-5) +self.assertAllClose(x_hat1, x_hat2) + +# Output values are in valid range for binary occupancy +self.assertAllGreaterEqual(output, 0.0) +self.assertAllLessEqual(output, 1.0) +``` + +**Bad assertions — do not write these:** +```python +self.assertIsNotNone(output) # Almost always true, tests nothing useful +self.assertTrue(isinstance(x, tf.Tensor)) # If it wasn't a tensor, the test already crashed +assert output.shape is not None # Tautological +self.assertEqual(len(output), len(output)) # Obviously true +assert model is not None # Construction already succeeded +``` + +**Note:** The existing test suite has some `assertIsNotNone` calls (e.g., in `test_model_transforms.py`). Do not add more of these. When modifying existing tests, replace them with meaningful assertions (shape checks, value range checks, dtype checks). + +### What Makes a Bad Test + +**Do not write tests that:** + +1. **Re-implement the function and check equality.** If your test duplicates the source logic to produce an expected value, it will always pass — even if both are wrong. + +2. **Mock the thing being tested.** Mock external dependencies, not the code under test. If you're testing `AnalysisTransform`, don't mock `Conv3D` — test with real convolutions on small tensors. + +3. **Only test the happy path.** Also test: empty batch (batch_size=0 if applicable), single-element batch, mismatched shapes, invalid entropy model string, out-of-range scale values. + +4. **Assert on random values without seeding.** If your test creates random tensors and asserts on specific values, set `tf.random.set_seed(42)` first. + +5. **Have side effects between test methods.** Each test method must be independent. Don't rely on execution order. Use fixtures, not class-level mutation. + +### Performance Tests + +Performance tests in `test_performance.py` assert timing bounds. When writing new ones: + +```python +@pytest.mark.slow +def test_optimized_is_faster(self): + # Time both approaches with enough iterations + start = time.perf_counter() + for _ in range(iterations): + _ = slow_approach() + slow_time = time.perf_counter() - start + + start = time.perf_counter() + for _ in range(iterations): + _ = fast_approach() + fast_time = time.perf_counter() - start + + speedup = slow_time / fast_time + # Use conservative thresholds — CI machines are slower than dev machines + assert speedup > 1.2, f"Expected >1.2x speedup, got {speedup:.1f}x" +``` + +Set conservative speedup thresholds (1.2x, not 10x) — CI environments vary. Always run multiple iterations to amortize startup. + +--- + +## Code Quality Standards + +### Two Contexts: Library Code vs CLI Tools + +The codebase has two distinct contexts with different conventions: + +**Library code** (`model_transforms.py`, `entropy_model.py`, `entropy_parameters.py`, `context_model.py`, `channel_context.py`, `attention_context.py`, `constants.py`, `precision_config.py`): +- No `print()` — these are imported modules, not user-facing scripts +- No `logging` — not used in library code +- Strict numerical discipline and `@tf.function` compatibility + +**CLI tools and pipeline scripts** (`benchmarks.py`, `quick_benchmark.py`, `cli_train.py`, `ds_*.py`, `ev_*.py`, `mp_*.py`, `training_pipeline.py`, `evaluation_pipeline.py`, `compress_octree.py`, `parallel_process.py`, `experiment.py`): +- `print()` is fine — these are user-facing programs that output to console +- `logging` is used and appropriate — several pipeline files use the `logging` module +- Standard Python error handling conventions apply + +When adding new code, follow the convention of the file you're editing. + +### Error Handling + +- **Do not add bare `except:` or `except Exception: pass` blocks** that silently swallow errors. TensorFlow errors (shape mismatches, OOM, NaN gradients) must propagate. There are a few existing `except Exception` blocks in `benchmarks.py` and `evaluation_pipeline.py` — do not add more. +- **Do not add try/except around tensor operations** in library code unless handling a specific, documented failure mode. +- Use `tf.debugging.assert_*` for development-time checks that can be disabled in production. + +### Linting & Imports + +- **Linter:** `ruff` (configured in `pyproject.toml`). No flake8, no pylint. +- Selected rule sets: **F** (Pyflakes), **I** (isort), **E/W** (pycodestyle). Do not add other rule sets without discussion. +- `sys.path.insert(0, ...)` is used in test files to find `src/`. Follow this pattern, don't fight it. +- Entropy model submodules use deferred imports inside `_create_entropy_model()` to avoid circular dependencies. Keep it this way. +- Import from `constants.py`, not inline math: `from constants import LOG_2_RECIPROCAL`. +- When adding new source modules, add the module name to `known-first-party` in `pyproject.toml` so isort classifies it correctly. + +### Typing + +- The codebase uses `typing` annotations (`Optional`, `Dict`, `Tuple`, etc.) in function signatures. +- Follow existing style. Don't add type annotations to existing functions that don't have them unless asked. + +### Things Not to Add Unless Asked + +- New dependencies in `requirements.txt` +- New base classes, ABCs, or mixin layers +- README updates to match code changes (the README describes some aspirational features) + +--- + +## Known Issues + +These are pre-existing issues in the codebase. Do not fix them unless specifically asked — they are documented here so you don't waste time rediscovering them: + +- `tests/test_point_cloud_metricss.py` — filename has a double 's' typo +- `tests/test_integration.py::TestIntegration` — entire class is `@pytest.mark.skip` due to API mismatch with `TrainingPipeline()` +- `pytest.ini` is missing the `slow` marker definition (tests use `@pytest.mark.slow` but it's not registered, causing a warning) +- Some test methods in `test_model_transforms.py` use `assertIsNotNone` where a shape or value assertion would be more meaningful +- Some test files (`test_data_loader`, `test_training_pipeline`, etc.) fail collection on machines with network timeouts — pre-existing, not caused by lint migration + +--- + +## Quick Reference + +```bash +# Lint (must pass before commit) +ruff check src/ tests/ + +# Auto-fix import order and whitespace +ruff check --fix src/ tests/ + +# Run all tests +pytest tests/ -v + +# Run one test file +pytest tests/test_entropy_model.py -v + +# Skip slow tests +pytest tests/ -v -m "not slow" + +# Run only GPU tests +pytest tests/ -v -m gpu + +# Quick smoke test (no data needed) +python -m src.quick_benchmark --compare +``` + +### Repository Layout + +``` +src/ + # Library code (strict quality rules) + model_transforms.py # AnalysisTransform, SynthesisTransform, DeepCompressModel(V2) + entropy_model.py # PatchedGaussianConditional, EntropyModel, MeanScaleHyperprior + entropy_parameters.py # Hyperprior μ/σ prediction network + context_model.py # MaskedConv3D, AutoregressiveContext + channel_context.py # ChannelContext, ChannelContextEntropyModel + attention_context.py # WindowedAttention3D, AttentionEntropyModel + constants.py # Pre-computed LOG_2, LOG_2_RECIPROCAL, EPSILON, etc. + precision_config.py # PrecisionManager for mixed float16 + + # CLI tools and pipeline scripts (standard Python conventions) + training_pipeline.py # TrainingPipeline class (uses logging) + evaluation_pipeline.py # EvaluationPipeline class (uses logging) + quick_benchmark.py # Synthetic smoke test (uses print) + benchmarks.py # Performance benchmarks (uses print) + cli_train.py # Training CLI entry point + ds_*.py # Data pipeline tools (mesh→pointcloud→octree→blocks) + ev_*.py, mp_*.py # Evaluation and MPEG comparison scripts + compress_octree.py # Compression entry point (uses logging) + octree_coding.py # Octree codec utility class + +tests/ + conftest.py # Session-scoped fixtures (tf_config, sample_point_cloud, file factories) + test_utils.py # Shared test utilities — USE THESE + test_*.py # One test file per source module +``` diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 000000000..0d00de1cd --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,27 @@ +[tool.ruff] +line-length = 120 +target-version = "py38" + +[tool.ruff.lint] +select = ["F", "I", "E", "W"] +ignore = [ + "E402", # module-level import not at top — tests use sys.path before imports + "E741", # ambiguous variable name — common in math-heavy code +] + +[tool.ruff.lint.per-file-ignores] +"tests/*.py" = ["E402"] + +[tool.ruff.lint.isort] +known-first-party = [ + "model_transforms", "entropy_model", "entropy_parameters", + "context_model", "channel_context", "attention_context", + "constants", "precision_config", "data_loader", + "training_pipeline", "evaluation_pipeline", "experiment", + "compress_octree", "decompress_octree", "octree_coding", + "ds_mesh_to_pc", "ds_pc_octree_blocks", "ds_select_largest", + "ev_compare", "ev_run_render", "mp_report", "mp_run", + "quick_benchmark", "benchmarks", "parallel_process", + "point_cloud_metrics", "map_color", "colorbar", + "cli_train", "test_utils", +] diff --git a/requirements.txt b/requirements.txt index 93bd3622f..035e50e82 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,3 +9,4 @@ pytest~=7.1.0 scipy~=1.8.1 numba~=0.56.0 tensorflow-probability~=0.19.0 +ruff>=0.4.0 diff --git a/src/attention_context.py b/src/attention_context.py index 11553e4b2..5cc6da477 100644 --- a/src/attention_context.py +++ b/src/attention_context.py @@ -11,8 +11,9 @@ - Global tokens provide long-range context without full attention """ +from typing import Any, Dict, Optional, Tuple + import tensorflow as tf -from typing import Tuple, Optional, Dict, Any from constants import LOG_2_RECIPROCAL @@ -668,8 +669,8 @@ def __init__(self, self.num_attention_layers = num_attention_layers # Import here to avoid circular dependency - from entropy_parameters import EntropyParameters from entropy_model import ConditionalGaussian + from entropy_parameters import EntropyParameters # Hyperprior-based parameter prediction self.entropy_parameters = EntropyParameters( @@ -803,9 +804,9 @@ def __init__(self, self.num_channel_groups = num_channel_groups self.num_attention_layers = num_attention_layers - from entropy_parameters import EntropyParameters from channel_context import ChannelContext from entropy_model import ConditionalGaussian + from entropy_parameters import EntropyParameters # Hyperprior parameters self.entropy_parameters = EntropyParameters( diff --git a/src/benchmarks.py b/src/benchmarks.py index bbe948474..709a47a5c 100644 --- a/src/benchmarks.py +++ b/src/benchmarks.py @@ -20,11 +20,12 @@ print(f"Peak memory: {mem.peak_mb:.1f} MB") """ -import tensorflow as tf import time -from typing import Callable, Dict, Any, Optional -from dataclasses import dataclass, field from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, Optional + +import tensorflow as tf @dataclass diff --git a/src/channel_context.py b/src/channel_context.py index 5bbd22241..60b1ba76f 100644 --- a/src/channel_context.py +++ b/src/channel_context.py @@ -7,8 +7,9 @@ maintaining autoregressive structure across groups. """ +from typing import Any, Dict, List, Optional, Tuple + import tensorflow as tf -from typing import Tuple, Optional, Dict, Any, List from constants import LOG_2_RECIPROCAL @@ -230,8 +231,8 @@ def __init__(self, self.channels_per_group = latent_channels // num_groups # Import here to avoid circular dependency - from entropy_parameters import EntropyParameters from entropy_model import ConditionalGaussian + from entropy_parameters import EntropyParameters # Hyperprior-based parameter prediction self.entropy_parameters = EntropyParameters( diff --git a/src/cli_train.py b/src/cli_train.py index 5d7a02c98..9eb21a00e 100644 --- a/src/cli_train.py +++ b/src/cli_train.py @@ -1,23 +1,25 @@ -import tensorflow as tf -import os import argparse import glob -import numpy as np +import os + import keras_tuner as kt +import tensorflow as tf + from ds_mesh_to_pc import read_off + def create_model(hp): model = tf.keras.Sequential() model.add(tf.keras.layers.InputLayer(input_shape=(2048, 3))) - + for i in range(hp.Int('num_layers', 1, 5)): model.add(tf.keras.layers.Dense( hp.Int(f'layer_{i}_units', min_value=64, max_value=1024, step=64), activation='relu' )) - + model.add(tf.keras.layers.Dense(3, activation='sigmoid')) - + model.compile( optimizer=tf.keras.optimizers.Adam( learning_rate=hp.Float('learning_rate', 1e-5, 1e-3, sampling='log') @@ -28,7 +30,7 @@ def create_model(hp): 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): mesh_data = read_off(file_path) return mesh_data.vertices @@ -47,7 +49,7 @@ def data_generator(): dataset = dataset.shuffle(buffer_size=len(file_paths)) dataset = dataset.batch(batch_size) dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE) - + return dataset def tune_hyperparameters(input_dir, output_dir, num_epochs=10): @@ -63,10 +65,10 @@ def tune_hyperparameters(input_dir, output_dir, num_epochs=10): dataset = load_and_preprocess_data(input_dir, batch_size=32) tuner.search(dataset, epochs=num_epochs, validation_data=dataset) - + best_model = tuner.get_best_models(num_models=1)[0] best_hps = tuner.get_best_hyperparameters(num_trials=1)[0] - + print("Best Hyperparameters:", best_hps.values) best_model.save(os.path.join(output_dir, 'best_model')) @@ -95,4 +97,4 @@ def main(): model.save(os.path.join(args.output_dir, 'trained_model')) if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/src/colorbar.py b/src/colorbar.py index 3f3734fbc..64301410a 100644 --- a/src/colorbar.py +++ b/src/colorbar.py @@ -1,10 +1,12 @@ -import numpy as np -import matplotlib.pyplot as plt -import matplotlib as mpl -from typing import Tuple, Callable, Optional, List -from dataclasses import dataclass import json +from dataclasses import dataclass from pathlib import Path +from typing import Callable, List, Optional, Tuple + +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np + @dataclass class ColorbarConfig: @@ -67,7 +69,10 @@ def get_colorbar( else: # Format numeric labels formatter = mpl.ticker.FormatStrFormatter(label_format) - cbar.ax.xaxis.set_major_formatter(formatter) if orientation == 'horizontal' else cbar.ax.yaxis.set_major_formatter(formatter) + if orientation == 'horizontal': + cbar.ax.xaxis.set_major_formatter(formatter) + else: + cbar.ax.yaxis.set_major_formatter(formatter) # Set font sizes cbar.ax.tick_params(labelsize=font_size) @@ -96,10 +101,10 @@ def save_color_mapping(filename: str, # Create mapping norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax) cmap = plt.get_cmap(cmap) - + values = np.linspace(vmin, vmax, num_samples) colors = [cmap(norm(v)) for v in values] - + # Save to file Path(filename).parent.mkdir(parents=True, exist_ok=True) with open(filename, 'w') as f: @@ -122,4 +127,4 @@ def save_color_mapping(filename: str, tick_rotation=45, extend='both' ) - plt.show() \ No newline at end of file + plt.show() diff --git a/src/compress_octree.py b/src/compress_octree.py index f6302e35b..dd738ac38 100644 --- a/src/compress_octree.py +++ b/src/compress_octree.py @@ -1,10 +1,11 @@ -import numpy as np -import os import logging -from typing import Tuple, Dict, Any, Optional, List -from pathlib import Path +import os +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np from scipy.spatial import cKDTree + class OctreeCompressor: def __init__(self, resolution: int = 64, debug_output: bool = False, output_dir: Optional[str] = None): self.resolution = resolution @@ -13,34 +14,36 @@ def __init__(self, resolution: int = 64, debug_output: bool = False, output_dir: if debug_output and output_dir: os.makedirs(output_dir, exist_ok=True) - def _create_voxel_grid(self, point_cloud: np.ndarray, normals: Optional[np.ndarray] = None) -> Tuple[np.ndarray, Dict[str, Any]]: + def _create_voxel_grid( + self, point_cloud: np.ndarray, normals: Optional[np.ndarray] = None, + ) -> Tuple[np.ndarray, Dict[str, Any]]: """Convert point cloud to voxel grid with enhanced metadata.""" grid = np.zeros((self.resolution,) * 3, dtype=bool) - + # Calculate bounds with epsilon to avoid division by zero min_bounds = np.min(point_cloud, axis=0) max_bounds = np.max(point_cloud, axis=0) ranges = max_bounds - min_bounds ranges = np.where(ranges == 0, 1e-6, ranges) - + # Scale points to grid resolution scaled_points = (point_cloud - min_bounds) / ranges * (self.resolution - 1) indices = np.clip(scaled_points, 0, self.resolution - 1).astype(int) - + # Mark occupied voxels for idx in indices: grid[tuple(idx)] = True - + metadata = { 'min_bounds': min_bounds, 'max_bounds': max_bounds, 'ranges': ranges, 'has_normals': normals is not None } - + if normals is not None: metadata['normal_grid'] = self._create_normal_grid(indices, normals) - + if self.debug_output and self.output_dir: debug_data = { 'grid': grid, @@ -48,89 +51,98 @@ def _create_voxel_grid(self, point_cloud: np.ndarray, normals: Optional[np.ndarr 'scaled_points': scaled_points } self._save_debug_info('grid_creation', debug_data) - + return grid, metadata def _create_normal_grid(self, indices: np.ndarray, normals: np.ndarray) -> np.ndarray: """Create grid storing average normals for occupied voxels.""" normal_grid = np.zeros((self.resolution, self.resolution, self.resolution, 3)) normal_counts = np.zeros((self.resolution, self.resolution, self.resolution)) - + # Accumulate normals in each voxel for idx, normal in zip(indices, normals): normal_grid[tuple(idx)] += normal normal_counts[tuple(idx)] += 1 - + # Average normals where counts > 0 with handling for zero counts counts_expanded = np.expand_dims(normal_counts, -1) counts_expanded = np.where(counts_expanded == 0, 1, counts_expanded) # Avoid division by zero normal_grid = normal_grid / counts_expanded - + # Normalize non-zero vectors norms = np.linalg.norm(normal_grid, axis=-1, keepdims=True) norms = np.where(norms == 0, 1, norms) # Avoid division by zero normal_grid = normal_grid / norms - + return normal_grid - def compress(self, point_cloud: np.ndarray, normals: Optional[np.ndarray] = None, validate: bool = True) -> Tuple[np.ndarray, Dict[str, Any]]: + 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") - + if normals is not None and normals.shape != point_cloud.shape: raise ValueError("Normals shape must match point cloud shape") - + grid, metadata = self._create_voxel_grid(point_cloud, normals) - + if validate: decompressed, _ = self.decompress(grid, metadata) error = self._compute_error(decompressed, point_cloud) metadata['compression_error'] = float(error) logging.info(f"Compression error: {error:.6f}") - + return grid, metadata - def decompress(self, grid: np.ndarray, metadata: Dict[str, Any], *, return_normals: bool = True) -> Tuple[np.ndarray, Optional[np.ndarray]]: + def decompress( + self, grid: np.ndarray, metadata: Dict[str, Any], *, + return_normals: bool = True, + ) -> Tuple[np.ndarray, Optional[np.ndarray]]: """ Decompress grid back to point cloud with optional normals. - + Args: grid: Compressed binary grid metadata: Compression metadata return_normals: Whether to return normals if available - + Returns: Tuple of points array and optional normals array """ indices = np.argwhere(grid).astype(np.float32) - + # Scale points back to original space points = indices / (self.resolution - 1) * metadata['ranges'] + metadata['min_bounds'] - + # Handle normals if present and requested normals = None if return_normals and metadata.get('has_normals') and 'normal_grid' in metadata: normals = metadata['normal_grid'][tuple(indices.astype(int).T)] - + return points, normals - 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]]]: + 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) - + def partition_recursive(points: np.ndarray, bounds: Tuple[np.ndarray, np.ndarray]) -> None: if len(points) <= max_points_per_block or np.min(bounds[1] - bounds[0]) <= min_block_size: if len(points) > 0: # Only add non-empty blocks blocks.append((points, {'bounds': bounds})) return - + mid = (bounds[0] + bounds[1]) / 2 for octant in np.ndindex((2, 2, 2)): # Compute octant bounds @@ -140,17 +152,17 @@ def partition_recursive(points: np.ndarray, bounds: Tuple[np.ndarray, np.ndarray max_corner = np.array([ mid[i] if octant[i] == 0 else bounds[1][i] for i in range(3) ]) - + # Find points in this octant with epsilon for stability epsilon = 1e-10 mask = np.all( - (points >= min_corner - epsilon) & + (points >= min_corner - epsilon) & (points <= max_corner + epsilon), axis=1 ) if np.any(mask): partition_recursive(points[mask], (min_corner, max_corner)) - + partition_recursive(point_cloud, (min_bounds, max_bounds)) return blocks @@ -164,10 +176,10 @@ def _save_debug_info(self, stage: str, data: Dict[str, Any]) -> None: """Save debug information to files.""" if not self.debug_output or not self.output_dir: return - + debug_dir = os.path.join(self.output_dir, 'debug', stage) os.makedirs(debug_dir, exist_ok=True) - + for name, array in data.items(): if isinstance(array, (np.ndarray, dict)): np.save(os.path.join(debug_dir, f"{name}.npy"), array) @@ -176,7 +188,7 @@ def save_compressed(self, grid: np.ndarray, metadata: Dict[str, Any], filename: """Save compressed data with metadata.""" os.makedirs(os.path.dirname(os.path.abspath(filename)), exist_ok=True) np.savez_compressed(filename, grid=grid, metadata=metadata) - + if self.debug_output: debug_path = f"{filename}.debug.npz" np.savez_compressed(debug_path, **metadata) @@ -184,4 +196,4 @@ def save_compressed(self, grid: np.ndarray, metadata: Dict[str, Any], filename: def load_compressed(self, filename: str) -> Tuple[np.ndarray, Dict[str, Any]]: """Load compressed data with metadata.""" data = np.load(filename, allow_pickle=True) - return data['grid'], data['metadata'].item() \ No newline at end of file + return data['grid'], data['metadata'].item() diff --git a/src/constants.py b/src/constants.py index ed898aa86..3ef3e40fd 100644 --- a/src/constants.py +++ b/src/constants.py @@ -11,9 +11,10 @@ bits = -likelihood * LOG_2_RECIPROCAL # Instead of / tf.math.log(2.0) """ -import tensorflow as tf import math +import tensorflow as tf + # Natural logarithm of 2: ln(2) = 0.693147... # Used for converting between natural log and log base 2 LOG_2 = tf.constant(math.log(2.0), dtype=tf.float32, name='log_2') diff --git a/src/context_model.py b/src/context_model.py index 68fb3265c..7fea40322 100644 --- a/src/context_model.py +++ b/src/context_model.py @@ -6,9 +6,10 @@ more accurate entropy estimation and better compression. """ -import tensorflow as tf +from typing import Any, Dict, Optional, Tuple + import numpy as np -from typing import Tuple, Optional, Dict, Any +import tensorflow as tf from constants import LOG_2_RECIPROCAL @@ -264,8 +265,8 @@ def __init__(self, self.num_context_layers = num_context_layers # Import here to avoid circular dependency - from entropy_parameters import EntropyParameters from entropy_model import ConditionalGaussian + from entropy_parameters import EntropyParameters # Hyperprior-based parameter prediction self.entropy_parameters = EntropyParameters( diff --git a/src/data_loader.py b/src/data_loader.py index c1dd4dd72..d48127d26 100644 --- a/src/data_loader.py +++ b/src/data_loader.py @@ -1,22 +1,24 @@ -import tensorflow as tf -import numpy as np import glob -import os from pathlib import Path -from typing import Dict, Any, Tuple, Optional +from typing import Any, Dict + +import numpy as np +import tensorflow as tf + from ds_mesh_to_pc import read_off from ds_pc_octree_blocks import PointCloudProcessor + class DataLoader: """Unified data loader for ModelNet40 and 8iVFB datasets.""" - + def __init__(self, config: Dict[str, Any]): self.config = config self.processor = PointCloudProcessor( block_size=config.get('block_size', 1.0), min_points=config.get('min_points', 100) ) - + def process_point_cloud(self, file_path: str) -> tf.Tensor: """Process a single point cloud file.""" # Read point cloud @@ -35,7 +37,7 @@ def process_point_cloud(self, file_path: str) -> tf.Tensor: voxelized = self._voxelize_points(points, resolution) return voxelized - + @tf.function def _normalize_points(self, points: tf.Tensor) -> tf.Tensor: """Normalize points to unit cube.""" @@ -44,9 +46,9 @@ def _normalize_points(self, points: tf.Tensor) -> tf.Tensor: scale = tf.reduce_max(tf.abs(points)) points = points / scale return points - + @tf.function - def _voxelize_points(self, + def _voxelize_points(self, points: tf.Tensor, resolution: int) -> tf.Tensor: """Convert points to voxel grid.""" @@ -54,12 +56,12 @@ def _voxelize_points(self, points = (points + 1) * (resolution - 1) / 2 points = tf.clip_by_value(points, 0, resolution - 1) indices = tf.cast(tf.round(points), tf.int32) - + # Create voxel grid grid = tf.zeros((resolution, resolution, resolution), dtype=tf.float32) updates = tf.ones(tf.shape(indices)[0], dtype=tf.float32) grid = tf.tensor_scatter_nd_update(grid, indices, updates) - + return grid def load_training_data(self) -> tf.data.Dataset: @@ -67,7 +69,7 @@ def load_training_data(self) -> tf.data.Dataset: train_path = Path(self.config['data']['modelnet40_path']) file_pattern = str(train_path / "**/*.off") files = glob.glob(file_pattern, recursive=True) - + # Create dataset dataset = tf.data.Dataset.from_tensor_slices(files) dataset = dataset.map( @@ -78,20 +80,20 @@ def load_training_data(self) -> tf.data.Dataset: ), num_parallel_calls=tf.data.AUTOTUNE ) - + # Batch and prefetch dataset = dataset.shuffle(1000) dataset = dataset.batch(self.config['training']['batch_size']) dataset = dataset.prefetch(tf.data.AUTOTUNE) - + return dataset - + def load_evaluation_data(self) -> tf.data.Dataset: """Load 8iVFB evaluation data.""" eval_path = Path(self.config['data']['ivfb_path']) file_pattern = str(eval_path / "*.ply") files = glob.glob(file_pattern) - + dataset = tf.data.Dataset.from_tensor_slices(files) dataset = dataset.map( lambda x: tf.py_function( @@ -101,13 +103,13 @@ def load_evaluation_data(self) -> tf.data.Dataset: ), num_parallel_calls=tf.data.AUTOTUNE ) - + # Batch without shuffling for evaluation dataset = dataset.batch(1) dataset = dataset.prefetch(tf.data.AUTOTUNE) - + return dataset - + @tf.function def _augment(self, points: tf.Tensor) -> tf.Tensor: """Apply data augmentation.""" @@ -119,7 +121,7 @@ def _augment(self, points: tf.Tensor) -> tf.Tensor: [0, 0, 1] ]) points = tf.matmul(points, rotation) - + # 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( @@ -128,4 +130,4 @@ def _augment(self, points: tf.Tensor) -> tf.Tensor: lambda: points ) - return points \ No newline at end of file + return points diff --git a/src/ds_mesh_to_pc.py b/src/ds_mesh_to_pc.py index 856fee5fe..72f19d0b8 100644 --- a/src/ds_mesh_to_pc.py +++ b/src/ds_mesh_to_pc.py @@ -1,9 +1,11 @@ -import numpy as np import argparse import logging -from pathlib import Path -from typing import Optional, Tuple, Dict, List from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +import numpy as np + @dataclass class MeshData: @@ -105,7 +107,7 @@ def sample_points_from_mesh( # Sample points indices = np.random.choice(len(areas), num_points, p=probabilities) points = np.array(centroids)[indices] - + # Get corresponding normals if requested normals = None if compute_normals and mesh_data.face_normals is not None: @@ -117,7 +119,7 @@ def sample_points_from_mesh( else: indices = np.random.choice(len(mesh_data.vertices), num_points, replace=True) points = mesh_data.vertices[indices] - + normals = None if compute_normals and mesh_data.vertex_normals is not None: normals = mesh_data.vertex_normals[indices] @@ -144,31 +146,27 @@ def partition_point_cloud( """ # Compute point cloud bounds min_bound = np.min(points, axis=0) - max_bound = np.max(points, axis=0) - - # Compute grid dimensions - grid_size = np.ceil((max_bound - min_bound) / block_size).astype(int) - + # Initialize blocks blocks = [] - + # Assign points to grid cells grid_indices = np.floor((points - min_bound) / block_size).astype(int) - + # Process each occupied grid cell unique_indices = np.unique(grid_indices, axis=0) for idx in unique_indices: # Find points in this cell mask = np.all(grid_indices == idx, axis=1) block_points = points[mask] - + # Only keep blocks with enough points if len(block_points) >= min_points: block_data = {'points': block_points} if normals is not None: block_data['normals'] = normals[mask] blocks.append(block_data) - + return blocks def save_ply( @@ -193,21 +191,21 @@ def save_ply( f.write("property float x\n") f.write("property float y\n") f.write("property float z\n") - + if normals is not None: f.write("property float nx\n") f.write("property float ny\n") f.write("property float nz\n") - + f.write("end_header\n") - + # Write data for i in range(len(points)): line = f"{points[i,0]} {points[i,1]} {points[i,2]}" if normals is not None: line += f" {normals[i,0]} {normals[i,1]} {normals[i,2]}" f.write(line + "\n") - + except Exception as e: logging.error(f"Error saving PLY file {file_path}: {e}") @@ -236,14 +234,14 @@ def convert_mesh_to_point_cloud( mesh_data = read_off(input_path) if mesh_data is None: return - + # Sample points points, normals = sample_points_from_mesh( mesh_data, num_points, compute_normals ) - + if partition_blocks: # Partition into blocks blocks = partition_point_cloud( @@ -252,7 +250,7 @@ def convert_mesh_to_point_cloud( block_size, min_points_per_block ) - + # Save each block output_base = Path(output_path).with_suffix('') for i, block in enumerate(blocks): @@ -284,12 +282,12 @@ def main(): help="Size of octree blocks (default: 1.0)") parser.add_argument("--min_points", type=int, default=100, help="Minimum points per block (default: 100)") - + args = parser.parse_args() - + # Configure logging logging.basicConfig(level=logging.INFO) - + # Convert mesh to point cloud convert_mesh_to_point_cloud( args.input, @@ -302,4 +300,4 @@ def main(): ) if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/src/ds_pc_octree_blocks.py b/src/ds_pc_octree_blocks.py index beb7c9317..d0791888f 100644 --- a/src/ds_pc_octree_blocks.py +++ b/src/ds_pc_octree_blocks.py @@ -1,12 +1,13 @@ -import tensorflow as tf import argparse -import os -from typing import List, Tuple from pathlib import Path +from typing import List + +import tensorflow as tf + class PointCloudProcessor: """Point cloud processing with TF 2.x operations.""" - + def __init__(self, block_size: float = 1.0, min_points: int = 10): self.block_size = block_size self.min_points = min_points @@ -16,11 +17,11 @@ def read_point_cloud(self, file_path: str) -> tf.Tensor: """Read point cloud using TF file operations.""" raw_data = tf.io.read_file(file_path) lines = tf.strings.split(raw_data, '\n')[1:] # Skip header - + def parse_line(line): values = tf.strings.split(line) return tf.strings.to_number(values[:3], out_type=tf.float32) - + points = tf.map_fn( parse_line, lines, @@ -32,42 +33,41 @@ def partition_point_cloud(self, points: tf.Tensor) -> List[tf.Tensor]: """Partition point cloud into blocks using TF operations.""" # Compute bounds min_bound = tf.reduce_min(points, axis=0) - max_bound = tf.reduce_max(points, axis=0) - + # Compute grid indices grid_indices = tf.cast( tf.floor((points - min_bound) / self.block_size), tf.int32 ) - + # Create unique block keys block_keys = ( grid_indices[:, 0] * 1000000 + grid_indices[:, 1] * 1000 + grid_indices[:, 2] ) - + # Get unique blocks unique_keys, inverse_indices = tf.unique(block_keys) - + blocks = [] for key in unique_keys: mask = tf.equal(block_keys, key) block_points = tf.boolean_mask(points, mask) - + if tf.shape(block_points)[0] >= self.min_points: blocks.append(block_points) - + return blocks def save_blocks(self, blocks: List[tf.Tensor], output_dir: str, base_name: str): """Save point cloud blocks to PLY files.""" output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) - + for i, block in enumerate(blocks): file_path = output_dir / f"{base_name}_block_{i}.ply" - + header = [ "ply", "format ascii 1.0", @@ -77,10 +77,10 @@ def save_blocks(self, blocks: List[tf.Tensor], output_dir: str, base_name: str): "property float z", "end_header" ] - + with open(file_path, 'w') as f: f.write('\n'.join(header) + '\n') - + # Convert points to strings and write points_str = tf.strings.reduce_join( tf.strings.as_string(block), @@ -98,7 +98,7 @@ def main(): description="Partition point clouds into octree blocks." ) parser.add_argument( - "input", + "input", type=str, help="Path to input PLY file" ) @@ -119,23 +119,23 @@ def main(): default=10, help="Minimum points per block" ) - + args = parser.parse_args() - + processor = PointCloudProcessor( block_size=args.block_size, min_points=args.min_points ) - + # Process point cloud points = processor.read_point_cloud(args.input) blocks = processor.partition_point_cloud(points) - + # Save blocks base_name = Path(args.input).stem processor.save_blocks(blocks, args.output_dir, base_name) - + print(f"Saved {len(blocks)} blocks to {args.output_dir}") if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/src/ds_select_largest.py b/src/ds_select_largest.py index e07d3de90..cae5a23e2 100644 --- a/src/ds_select_largest.py +++ b/src/ds_select_largest.py @@ -1,7 +1,9 @@ -import os import argparse +import os + import tensorflow as tf + def count_points_in_block(file_path): """ Counts the number of points in a block from a `.ply` file. diff --git a/src/entropy_model.py b/src/entropy_model.py index a23c5a18b..47bcbd185 100644 --- a/src/entropy_model.py +++ b/src/entropy_model.py @@ -1,6 +1,7 @@ +from typing import Any, Dict, Optional, Tuple + import tensorflow as tf import tensorflow_probability as tfp -from typing import Optional, Dict, Any, Tuple from constants import LOG_2_RECIPROCAL diff --git a/src/entropy_parameters.py b/src/entropy_parameters.py index 1120ece97..2da78d867 100644 --- a/src/entropy_parameters.py +++ b/src/entropy_parameters.py @@ -6,8 +6,9 @@ where the distribution is adapted based on learned side information. """ +from typing import Optional, Tuple + import tensorflow as tf -from typing import Tuple, Optional class EntropyParameters(tf.keras.layers.Layer): diff --git a/src/ev_compare.py b/src/ev_compare.py index 17708be8b..6f79dad60 100644 --- a/src/ev_compare.py +++ b/src/ev_compare.py @@ -1,11 +1,11 @@ -import tensorflow as tf -import os import argparse -import time import logging -from typing import Dict, Tuple -from pathlib import Path +import os from dataclasses import dataclass +from typing import Dict + +import tensorflow as tf + @dataclass class EvaluationConfig: @@ -16,7 +16,7 @@ class EvaluationConfig: class PointCloudMetrics(tf.keras.metrics.Metric): """Custom metric class for point cloud evaluation.""" - + def __init__(self, name='point_cloud_metrics', **kwargs): super().__init__(name=name, **kwargs) self.psnr = self.add_weight(name='psnr', initializer='zeros') @@ -31,23 +31,23 @@ def compute_psnr(self, original: tf.Tensor, compressed: tf.Tensor) -> tf.Tensor: return 10 * tf.math.log(max_val**2 / mse) / tf.math.log(10.0) @tf.function - def compute_chamfer(self, + def compute_chamfer(self, original: tf.Tensor, compressed: tf.Tensor) -> tf.Tensor: """Compute Chamfer distance.""" # Compute pairwise distances original_expanded = tf.expand_dims(original, 1) compressed_expanded = tf.expand_dims(compressed, 0) - + distances = tf.reduce_sum( tf.square(original_expanded - compressed_expanded), axis=-1 ) - + # Compute minimum distances in both directions d1 = tf.reduce_min(distances, axis=1) d2 = tf.reduce_min(distances, axis=0) - + return tf.reduce_mean(d1) + tf.reduce_mean(d2) @tf.function @@ -55,7 +55,7 @@ def update_state(self, original: tf.Tensor, compressed: tf.Tensor): """Update metric states.""" psnr = self.compute_psnr(original, compressed) chamfer = self.compute_chamfer(original, compressed) - + self.psnr.assign_add(psnr) self.chamfer.assign_add(chamfer) self.count.assign_add(1) @@ -75,28 +75,28 @@ def reset_state(self): class CompressionEvaluator: """Evaluator for point cloud compression.""" - + def __init__(self, config: EvaluationConfig): self.config = config self.metrics = PointCloudMetrics() - + @tf.function def load_point_cloud(self, file_path: str) -> tf.Tensor: """Load point cloud from file.""" raw_data = tf.io.read_file(file_path) lines = tf.strings.split(raw_data, '\n') - + # Skip header data_lines = lines[tf.where( tf.strings.regex_full_match(lines, r'[\d\.\-\+eE\s]+') )[:, 0]] - + # Parse points points = tf.strings.to_number( tf.strings.split(data_lines), out_type=tf.float32 ) - + return points[:, :3] # Return only XYZ coordinates def evaluate_compression(self, @@ -106,19 +106,19 @@ def evaluate_compression(self, # Load point clouds original = self.load_point_cloud(original_path) compressed = self.load_point_cloud(compressed_path) - + # Update metrics self.metrics.update_state(original, compressed) - + # Get results results = self.metrics.result() - + # Add additional metrics results['file_size_ratio'] = ( os.path.getsize(compressed_path) / os.path.getsize(original_path) ) - + return {k: float(v) for k, v in results.items()} def main(): @@ -141,39 +141,39 @@ def main(): default="evaluation_results.json", help="Path to save evaluation results" ) - + args = parser.parse_args() - + # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) - + # Initialize evaluator evaluator = CompressionEvaluator(EvaluationConfig()) - + # Find matching files original_files = set(os.listdir(args.original_dir)) compressed_files = set(os.listdir(args.compressed_dir)) common_files = original_files & compressed_files - + results = {} for filename in common_files: logger.info(f"Evaluating {filename}") - + original_path = os.path.join(args.original_dir, filename) compressed_path = os.path.join(args.compressed_dir, filename) - + results[filename] = evaluator.evaluate_compression( original_path, compressed_path ) - + # Save results import json with open(args.output, 'w') as f: json.dump(results, f, indent=2) - + logger.info(f"Results saved to {args.output}") if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/src/ev_run_render.py b/src/ev_run_render.py index 137a2008e..c783d3d8f 100644 --- a/src/ev_run_render.py +++ b/src/ev_run_render.py @@ -2,12 +2,14 @@ import json import logging import os -from typing import Dict, Any, Tuple, Optional -import yaml -import tensorflow as tf +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple + import numpy as np +import tensorflow as tf +import yaml from PIL import Image -from dataclasses import dataclass + @dataclass class RenderConfig: @@ -34,30 +36,30 @@ class CameraParams: class PointCloudRenderer: """Point cloud renderer using TensorFlow.""" - + def __init__(self, config: RenderConfig): self.config = config self._setup_colormap() - + def _setup_colormap(self): """Setup color mapping for point cloud visualization.""" import matplotlib.pyplot as plt self.colormap = plt.get_cmap(self.config.color_map) - + def _compute_projection_matrix(self, camera: CameraParams) -> tf.Tensor: """Compute perspective projection matrix.""" f = 1.0 / tf.tan(tf.constant(camera.fov * 0.5 * np.pi / 180.0, dtype=tf.float32)) aspect = tf.constant(camera.aspect, dtype=tf.float32) near = tf.constant(camera.near, dtype=tf.float32) far = tf.constant(camera.far, dtype=tf.float32) - + return tf.cast(tf.stack([ [f/aspect, 0, 0, 0], [0, f, 0, 0], [0, 0, (far+near)/(near-far), 2*far*near/(near-far)], [0, 0, -1, 0] ]), tf.float32) - + def _compute_view_matrix(self, camera: CameraParams) -> tf.Tensor: """Compute view matrix from camera parameters.""" position = tf.cast(camera.position, tf.float32) @@ -72,10 +74,10 @@ def _compute_view_matrix(self, camera: CameraParams) -> tf.Tensor: # Compute camera axes z_axis = position - target z_axis = tf.nn.l2_normalize(z_axis, axis=-1) - + x_axis = tf.linalg.cross(up, z_axis) x_axis = tf.nn.l2_normalize(x_axis, axis=-1) - + y_axis = tf.linalg.cross(z_axis, x_axis) # Remove extra dimensions @@ -90,7 +92,7 @@ def _compute_view_matrix(self, camera: CameraParams) -> tf.Tensor: -tf.reduce_sum(y_axis * position), -tf.reduce_sum(z_axis * position) ]) - + # Build view matrix view_matrix = tf.stack([ tf.concat([x_axis, [translation[0]]], axis=0), @@ -98,45 +100,45 @@ def _compute_view_matrix(self, camera: CameraParams) -> tf.Tensor: tf.concat([z_axis, [translation[2]]], axis=0), tf.constant([0.0, 0.0, 0.0, 1.0], dtype=tf.float32) ]) - + return view_matrix - + def _project_points(self, points: tf.Tensor, camera: CameraParams) -> tf.Tensor: """Project 3D points to 2D image coordinates.""" # Ensure points are float32 points = tf.cast(points, tf.float32) - + # Homogeneous coordinates points_h = tf.concat([points, tf.ones([tf.shape(points)[0], 1], dtype=tf.float32)], axis=1) - + # View and projection transformations view_matrix = self._compute_view_matrix(camera) proj_matrix = self._compute_projection_matrix(camera) - + # Transform points view_points = tf.matmul(points_h, view_matrix, transpose_b=True) proj_points = tf.matmul(view_points, proj_matrix, transpose_b=True) - + # Perspective division proj_points = proj_points[..., :2] / (proj_points[..., 3:4] + 1e-10) - + # Scale to image coordinates image_points = (proj_points + 1.0) * 0.5 image_points = tf.stack([ image_points[:, 0] * tf.cast(self.config.image_width, tf.float32), (1.0 - image_points[:, 1]) * tf.cast(self.config.image_height, tf.float32) ], axis=1) - + return image_points - - def render(self, - points: tf.Tensor, - colors: Optional[tf.Tensor] = None, + + def render(self, + points: tf.Tensor, + colors: Optional[tf.Tensor] = None, camera: Optional[CameraParams] = None) -> Tuple[np.ndarray, Dict[str, Any]]: """Render point cloud to image.""" # Convert points to float32 points = tf.cast(points, tf.float32) - + if camera is None: # Compute default camera parameters center = tf.reduce_mean(points, axis=0) @@ -148,23 +150,23 @@ def render(self, fov=45.0, aspect=self.config.image_width / self.config.image_height ) - + # Project points to 2D image_points = self._project_points(points, camera) - + # Initialize image image = tf.zeros((self.config.image_height, self.config.image_width, 3), dtype=tf.float32) - + # Get point colors if colors is None: # Color by depth depths = tf.norm(points - tf.cast(camera.position, tf.float32), axis=1) norm_depths = (depths - tf.reduce_min(depths)) / (tf.reduce_max(depths) - tf.reduce_min(depths) + 1e-10) colors = tf.convert_to_tensor([self.colormap(d) for d in norm_depths.numpy()])[:, :3] - + # Convert colors to float32 colors = tf.cast(colors, tf.float32) - + # Render points valid_mask = tf.logical_and( tf.logical_and( @@ -176,20 +178,20 @@ def render(self, image_points[:, 1] < self.config.image_height ) ) - + valid_points = tf.boolean_mask(image_points, valid_mask) valid_colors = tf.boolean_mask(colors, valid_mask) - + # Simple point splatting indices = tf.cast(tf.round(valid_points), tf.int32) updates = valid_colors - + image = tf.tensor_scatter_nd_update( image, indices, updates ) - + render_info = { 'camera': { 'position': camera.position.tolist(), @@ -207,24 +209,24 @@ def render(self, 'color_map': self.config.color_map } } - + return image.numpy(), render_info def load_experiment_config(experiment_path: str) -> Dict[str, Any]: """Load and validate experiment configuration.""" with open(experiment_path, 'r') as f: config = yaml.safe_load(f) - + required_keys = [ - 'MPEG_DATASET_DIR', - 'EXPERIMENT_DIR', - 'model_configs', + 'MPEG_DATASET_DIR', + 'EXPERIMENT_DIR', + 'model_configs', 'vis_comps' ] missing_keys = [key for key in required_keys if key not in config] if missing_keys: raise ValueError(f"Missing required keys in configuration: {missing_keys}") - + return config def save_rendered_image( @@ -235,13 +237,13 @@ def save_rendered_image( ): """Save rendered image with metadata.""" img = Image.fromarray((image_array * 255).astype(np.uint8)) - + if bbox is not None: img = img.crop(bbox) render_info['bbox'] = bbox - + img.save(save_path) - + with open(save_path + ".meta.json", 'w') as f: json.dump(render_info, f, indent=2) @@ -253,31 +255,31 @@ def main(experiment_path: str): datefmt="%Y-%m-%d %H:%M:%S" ) logger = logging.getLogger(__name__) - + config = load_experiment_config(experiment_path) renderer = PointCloudRenderer(RenderConfig(**config.get('render_config', {}))) - + for data_entry in config['data']: pc_name = data_entry['pc_name'] output_dir = os.path.join(config['EXPERIMENT_DIR'], pc_name) os.makedirs(output_dir, exist_ok=True) - + point_cloud = tf.random.uniform((1024, 3), dtype=tf.float32) - + camera = None if 'camera_params' in data_entry: camera = CameraParams(**data_entry['camera_params']) - + image, render_info = renderer.render(point_cloud, camera=camera) bbox = data_entry.get('bbox', None) - + save_path = os.path.join(output_dir, f"{pc_name}.png") save_rendered_image(image, render_info, save_path, bbox) - + logger.info(f"Rendered {pc_name} and saved at {save_path}") if __name__ == '__main__': parser = argparse.ArgumentParser(description="Run point cloud rendering experiments.") parser.add_argument('experiment_path', help="Path to experiment configuration YAML.") args = parser.parse_args() - main(args.experiment_path) \ No newline at end of file + main(args.experiment_path) diff --git a/src/evaluation_pipeline.py b/src/evaluation_pipeline.py index 105c8cf72..0a44e5fc8 100644 --- a/src/evaluation_pipeline.py +++ b/src/evaluation_pipeline.py @@ -1,15 +1,16 @@ -import tensorflow as tf -import os import logging +from dataclasses import asdict, dataclass from pathlib import Path -from typing import Dict, Any, List -from dataclasses import dataclass, asdict +from typing import Any, Dict, List + +import tensorflow as tf from data_loader import DataLoader -from model_transforms import DeepCompressModel, TransformConfig from ev_compare import PointCloudMetrics +from model_transforms import DeepCompressModel, TransformConfig from mp_report import ExperimentReporter + @dataclass class EvaluationResult: """Container for evaluation results.""" @@ -22,22 +23,22 @@ class EvaluationResult: class EvaluationPipeline: """Pipeline for evaluating DeepCompress model.""" - + def __init__(self, config_path: str): self.config = self._load_config(config_path) self.logger = logging.getLogger(__name__) - + # Initialize components self.data_loader = DataLoader(self.config) self.metrics = PointCloudMetrics() self.model = self._load_model() - + def _load_config(self, config_path: str) -> Dict[str, Any]: """Load configuration from YAML.""" import yaml with open(config_path, 'r') as f: return yaml.safe_load(f) - + def _load_model(self) -> DeepCompressModel: """Load model from checkpoint.""" model_config = TransformConfig( @@ -45,18 +46,18 @@ def _load_model(self) -> DeepCompressModel: activation=self.config['model'].get('activation', 'cenic_gdn'), conv_type=self.config['model'].get('conv_type', 'separable') ) - + model = DeepCompressModel(model_config) - + # Load weights if checkpoint provided checkpoint_path = self.config.get('checkpoint_path') if checkpoint_path: model.load_weights(checkpoint_path) - + return model - + def _evaluate_single(self, - point_cloud: tf.Tensor) -> Dict[str, tf.Tensor]: + point_cloud) -> Dict[str, float]: """Evaluate model on single point cloud.""" # Forward pass through model x_hat, y, y_hat, z = self.model(point_cloud, training=False) @@ -67,12 +68,12 @@ def _evaluate_single(self, results['chamfer'] = self.metrics.compute_chamfer(point_cloud, x_hat) return results - - def evaluate(self) -> Dict[str, List[EvaluationResult]]: + + def evaluate(self) -> Dict[str, EvaluationResult]: """Run evaluation on test dataset.""" results = {} dataset = self.data_loader.load_evaluation_data() - + for i, point_cloud in enumerate(dataset): filename = f"point_cloud_{i}" self.logger.info(f"Evaluating {filename}") @@ -98,9 +99,9 @@ def evaluate(self) -> Dict[str, List[EvaluationResult]]: except Exception as e: self.logger.error(f"Error processing {filename}: {str(e)}") continue - + return results - + def generate_report(self, results: Dict[str, EvaluationResult]): """Generate evaluation report.""" # Convert EvaluationResult objects to flat dicts for ExperimentReporter @@ -111,33 +112,33 @@ def generate_report(self, results: Dict[str, EvaluationResult]): else: flat_results[name] = result reporter = ExperimentReporter(flat_results) - + # Generate and save report output_dir = Path(self.config['evaluation']['output_dir']) output_dir.mkdir(parents=True, exist_ok=True) - + report_path = output_dir / "evaluation_report.json" reporter.save_report(str(report_path)) - + self.logger.info(f"Evaluation report saved to {report_path}") - + def main(): import argparse parser = argparse.ArgumentParser(description="Evaluate DeepCompress model") parser.add_argument("config", type=str, help="Path to config file") parser.add_argument("--checkpoint", type=str, help="Path to model checkpoint") args = parser.parse_args() - + # Configure logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) - + # Run evaluation pipeline = EvaluationPipeline(args.config) results = pipeline.evaluate() pipeline.generate_report(results) if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/src/experiment.py b/src/experiment.py index d4ba626b0..481fc3228 100644 --- a/src/experiment.py +++ b/src/experiment.py @@ -1,7 +1,9 @@ -import yaml -import os import logging -from typing import Dict, Any +import os +from typing import Any, Dict + +import yaml + class ExperimentConfig: def __init__(self, config_path: str): @@ -57,4 +59,4 @@ def _run_model_pipeline(self, model_config: Dict[str, Any]): parser.add_argument("config", type=str, help="Path to the experiment configuration YAML file.") args = parser.parse_args() experiment = Experiment(args.config) - experiment.run() \ No newline at end of file + experiment.run() diff --git a/src/map_color.py b/src/map_color.py index 5da9962dd..5acd98c34 100644 --- a/src/map_color.py +++ b/src/map_color.py @@ -1,8 +1,9 @@ -import numpy as np import argparse -from pathlib import Path from typing import Optional +import numpy as np + + def load_point_cloud(file_path: str) -> Optional[np.ndarray]: """ Load a point cloud from a PLY file. diff --git a/src/model_transforms.py b/src/model_transforms.py index b13f7b190..aed2c2bb2 100644 --- a/src/model_transforms.py +++ b/src/model_transforms.py @@ -1,6 +1,7 @@ -import tensorflow as tf -from typing import Tuple from dataclasses import dataclass +from typing import Tuple + +import tensorflow as tf from constants import LOG_2_RECIPROCAL, EPSILON diff --git a/src/mp_report.py b/src/mp_report.py index 1a7e3ffce..b2e6aa22b 100644 --- a/src/mp_report.py +++ b/src/mp_report.py @@ -1,10 +1,11 @@ -import tensorflow as tf import json -import os -from typing import Dict, Any, List -from pathlib import Path -from dataclasses import dataclass import logging +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict + +import tensorflow as tf + @dataclass class ExperimentMetrics: @@ -18,12 +19,12 @@ class ExperimentMetrics: class ExperimentReporter: """Reporter for compression experiments.""" - + def __init__(self, experiment_results: Dict[str, Any]): self.experiment_results = experiment_results self.summary = self._initialize_summary() self.logger = logging.getLogger(__name__) - + def _initialize_summary(self) -> Dict[str, Any]: """Initialize summary metrics.""" return { @@ -34,34 +35,34 @@ def _initialize_summary(self) -> Dict[str, Any]: 'timestamp': self.experiment_results.get('timestamp', 'N/A') } } - + @tf.function def compute_aggregate_metrics(self) -> Dict[str, tf.Tensor]: """Compute aggregate metrics using TensorFlow operations.""" metrics = [] - + for file_name, results in self.experiment_results.items(): if file_name in ['timestamp', 'octree_levels', 'quantization_levels']: continue - + if all(key in results for key in ['psnr', 'bd_rate', 'bitrate']): metrics.append([ results['psnr'], results['bd_rate'], results['bitrate'] ]) - + if not metrics: return {} - + metrics_tensor = tf.convert_to_tensor(metrics, dtype=tf.float32) - + return { 'avg_psnr': tf.reduce_mean(metrics_tensor[:, 0]), 'avg_bd_rate': tf.reduce_mean(metrics_tensor[:, 1]), 'avg_bitrate': tf.reduce_mean(metrics_tensor[:, 2]) } - + def _compute_best_metrics(self) -> Dict[str, Any]: """Compute best metrics across all experiments.""" best_metrics = { @@ -72,7 +73,7 @@ def _compute_best_metrics(self) -> Dict[str, Any]: 'compression_time': float('inf'), 'decompression_time': float('inf') } - + best_models = { 'psnr': None, 'bd_rate': None, @@ -81,11 +82,11 @@ def _compute_best_metrics(self) -> Dict[str, Any]: 'compression_time': None, 'decompression_time': None } - + for file_name, results in self.experiment_results.items(): if file_name in ['timestamp', 'octree_levels', 'quantization_levels']: continue - + # Update best metrics for metric in best_metrics.keys(): if metric in results: @@ -98,7 +99,7 @@ def _compute_best_metrics(self) -> Dict[str, Any]: if value < best_metrics[metric]: best_metrics[metric] = value best_models[metric] = file_name - + return { 'metrics': best_metrics, 'models': best_models @@ -170,10 +171,10 @@ def save_report(self, output_file: str): report = self.generate_report() output_path = Path(output_file) output_path.parent.mkdir(parents=True, exist_ok=True) - + with open(output_path, 'w') as f: json.dump(report, f, indent=4) - + self.logger.info(f"Report saved to {output_file}") def load_experiment_results(input_file: str) -> Dict[str, Any]: @@ -188,7 +189,7 @@ def main(): format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) logger = logging.getLogger(__name__) - + # Parse arguments import argparse parser = argparse.ArgumentParser(description="Generate experiment report") @@ -203,16 +204,16 @@ def main(): help="Path to save the generated report" ) args = parser.parse_args() - + try: # Load results and generate report results = load_experiment_results(args.input_file) reporter = ExperimentReporter(results) reporter.save_report(args.output_file) - + except Exception as e: logger.error(f"Error generating report: {str(e)}", exc_info=True) raise if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/src/octree_coding.py b/src/octree_coding.py index 6c268b818..1e9fc33b9 100644 --- a/src/octree_coding.py +++ b/src/octree_coding.py @@ -1,6 +1,8 @@ -import tensorflow as tf -from typing import List, Tuple, Dict, Any from dataclasses import dataclass +from typing import Any, Dict, List, Tuple + +import tensorflow as tf + @dataclass class OctreeConfig: @@ -11,11 +13,11 @@ class OctreeConfig: class OctreeCoder(tf.keras.layers.Layer): """TensorFlow 2.x implementation of octree coding.""" - + def __init__(self, config: OctreeConfig, **kwargs): super().__init__(**kwargs) self.config = config - + def encode(self, point_cloud: tf.Tensor) -> Tuple[tf.Tensor, Dict[str, Any]]: """Encode point cloud into octree representation.""" # Create empty grid @@ -23,19 +25,19 @@ def encode(self, point_cloud: tf.Tensor) -> Tuple[tf.Tensor, Dict[str, Any]]: (self.config.resolution,) * 3, dtype=tf.bool ) - + # Calculate bounds min_bounds = tf.reduce_min(point_cloud, axis=0) max_bounds = tf.reduce_max(point_cloud, axis=0) scale = max_bounds - min_bounds - + # Handle zero scales scale = tf.where( tf.equal(scale, 0), tf.ones_like(scale) * self.config.epsilon, scale ) - + # Scale points to grid resolution scaled_points = (point_cloud - min_bounds) / scale * tf.cast( self.config.resolution - 1, @@ -49,21 +51,21 @@ def encode(self, point_cloud: tf.Tensor) -> Tuple[tf.Tensor, Dict[str, Any]]: ), tf.int32 ) - + # Create update values updates = tf.ones(tf.shape(indices)[0], dtype=tf.bool) - + # Update grid grid = tf.tensor_scatter_nd_update(grid, indices, updates) - + metadata = { 'min_bounds': min_bounds.numpy(), 'max_bounds': max_bounds.numpy(), 'scale': scale.numpy() } - + return grid, metadata - + def decode(self, grid: tf.Tensor, metadata: Dict[str, Any]) -> tf.Tensor: @@ -73,17 +75,17 @@ def decode(self, tf.where(grid), tf.float32 ) - + # Scale back to original space scale = tf.constant(metadata['scale'], dtype=tf.float32) min_bounds = tf.constant(metadata['min_bounds'], dtype=tf.float32) - + points = ( indices / tf.cast(self.config.resolution - 1, tf.float32) ) * scale + min_bounds - + return points - + def partition_octree( self, point_cloud: tf.Tensor, @@ -93,12 +95,12 @@ def partition_octree( """Partition point cloud into octree blocks.""" if level == 0 or point_cloud.shape[0] == 0: return [(point_cloud, bbox)] - + xmin, xmax, ymin, ymax, zmin, zmax = bbox xmid = (xmin + xmax) / 2 ymid = (ymin + ymax) / 2 zmid = (zmin + zmax) / 2 - + blocks = [] ranges = [ ((xmin, xmid), (ymin, ymid), (zmin, zmid)), @@ -110,7 +112,7 @@ def partition_octree( ((xmid, xmax), (ymid, ymax), (zmin, zmid)), ((xmid, xmax), (ymid, ymax), (zmid, zmax)) ] - + for x_range, y_range, z_range in ranges: # Compute conditions x_cond = tf.logical_and( @@ -125,13 +127,13 @@ def partition_octree( point_cloud[:, 2] >= z_range[0] - self.config.epsilon, point_cloud[:, 2] <= z_range[1] + self.config.epsilon ) - + # Combine conditions mask = tf.logical_and(x_cond, tf.logical_and(y_cond, z_cond)) - + # Get points in block in_block = tf.boolean_mask(point_cloud, mask) - + if in_block.shape[0] > 0: child_bbox = ( x_range[0], x_range[1], @@ -145,5 +147,5 @@ def partition_octree( level - 1 ) ) - - return blocks \ No newline at end of file + + return blocks diff --git a/src/parallel_process.py b/src/parallel_process.py index 9b4d846a9..135d74189 100644 --- a/src/parallel_process.py +++ b/src/parallel_process.py @@ -1,11 +1,11 @@ -import multiprocessing -import subprocess import logging +import subprocess import time -from typing import Callable, Any, List, Dict, Optional, Union from concurrent.futures import ThreadPoolExecutor, TimeoutError from dataclasses import dataclass from queue import Queue +from typing import Any, Callable, Dict, List, Optional + @dataclass class ProcessResult: @@ -20,9 +20,9 @@ class ProcessTimeoutError(Exception): pass class Popen: - def __init__(self, - cmd: List[str], - stdout: Any = None, + def __init__(self, + cmd: List[str], + stdout: Any = None, stderr: Any = None, timeout: Optional[float] = None): """ @@ -31,7 +31,7 @@ def __init__(self, self.cmd = cmd self.timeout = timeout self.start_time = time.time() - + self.process = subprocess.Popen( cmd, stdout=stdout, @@ -42,7 +42,7 @@ def __init__(self, def wait(self, timeout: Optional[float] = None) -> int: """Wait for the subprocess to complete with timeout.""" wait_timeout = timeout or self.timeout - + if wait_timeout is not None: try: return self.process.wait(timeout=wait_timeout) @@ -51,7 +51,7 @@ def wait(self, timeout: Optional[float] = None) -> int: raise ProcessTimeoutError( f"Process timed out after {wait_timeout} seconds: {' '.join(self.cmd)}" ) - + return self.process.wait() def terminate(self): @@ -121,22 +121,22 @@ def worker(index: int, param: Any, result_queue: Queue) -> None: # Process all parameters in parallel result_queue: Queue = Queue() results_dict: Dict[int, ProcessResult] = {} - + with ThreadPoolExecutor(max_workers=num_parallel) as executor: futures = [ executor.submit(worker, i, param, result_queue) for i, param in enumerate(params_list) ] - + # Wait for all tasks to complete for future in futures: future.result() # This will propagate any exceptions - + # Collect results and maintain order while len(results_dict) < len(params_list): result = result_queue.get() results_dict[result.index] = result - + # Process results in order ordered_results = [] for i in range(len(params_list)): @@ -146,5 +146,5 @@ def worker(index: int, param: Any, result_queue: Queue) -> None: raise TimeoutError(str(result.error)) raise result.error or RuntimeError(f"Task {i} failed without specific error") ordered_results.append(result.result) - - return ordered_results \ No newline at end of file + + return ordered_results diff --git a/src/point_cloud_metrics.py b/src/point_cloud_metrics.py index ee6ca75ee..dad73ef18 100644 --- a/src/point_cloud_metrics.py +++ b/src/point_cloud_metrics.py @@ -1,15 +1,17 @@ +from typing import Dict, Optional + import numpy as np from numba import njit, prange from scipy.spatial import cKDTree -from typing import Tuple, Optional, Dict + @njit(parallel=True) -def compute_point_to_point_distances(points1: np.ndarray, +def compute_point_to_point_distances(points1: np.ndarray, points2: np.ndarray) -> np.ndarray: N = points1.shape[0] M = points2.shape[0] distances = np.empty(N, dtype=np.float32) - + for i in prange(N): min_dist = np.inf for j in range(M): @@ -17,17 +19,17 @@ def compute_point_to_point_distances(points1: np.ndarray, if dist < min_dist: min_dist = dist distances[i] = np.sqrt(min_dist) - + return distances @njit(parallel=True) -def compute_point_to_normal_distances(points1: np.ndarray, +def compute_point_to_normal_distances(points1: np.ndarray, points2: np.ndarray, normals2: np.ndarray) -> np.ndarray: N = points1.shape[0] M = points2.shape[0] distances = np.empty(N, dtype=np.float32) - + for i in prange(N): min_dist = np.inf for j in range(M): @@ -36,7 +38,7 @@ def compute_point_to_normal_distances(points1: np.ndarray, if dist < min_dist: min_dist = dist distances[i] = min_dist - + return distances def calculate_metrics(predicted: np.ndarray, @@ -54,7 +56,7 @@ def calculate_metrics(predicted: np.ndarray, raise ValueError("Empty point cloud provided") if predicted.shape[1] != 3 or ground_truth.shape[1] != 3: raise ValueError("Point clouds must have shape (N, 3)") - + metrics = {} if use_kdtree: tree_gt = cKDTree(ground_truth) @@ -64,11 +66,11 @@ def calculate_metrics(predicted: np.ndarray, else: d1_distances = compute_point_to_point_distances(predicted, ground_truth) d2_distances = compute_point_to_point_distances(ground_truth, predicted) - + metrics['d1'] = np.mean(d1_distances) metrics['d2'] = np.mean(d2_distances) metrics['chamfer'] = metrics['d1'] + metrics['d2'] - + if predicted_normals is not None and ground_truth_normals is not None: if use_kdtree: _, indices_gt = tree_gt.query(predicted, k=1) @@ -87,7 +89,7 @@ def calculate_metrics(predicted: np.ndarray, metrics['n1'] = np.mean(n1_distances) metrics['n2'] = np.mean(n2_distances) metrics['normal_chamfer'] = metrics['n1'] + metrics['n2'] - + return metrics def calculate_chamfer_distance(predicted: np.ndarray, target: np.ndarray) -> float: @@ -96,4 +98,4 @@ def calculate_chamfer_distance(predicted: np.ndarray, target: np.ndarray) -> flo def calculate_d1_metric(predicted: np.ndarray, target: np.ndarray) -> float: metrics = calculate_metrics(predicted, target) - return metrics["d1"] \ No newline at end of file + return metrics["d1"] diff --git a/src/precision_config.py b/src/precision_config.py index 80caa5794..5db611db5 100644 --- a/src/precision_config.py +++ b/src/precision_config.py @@ -18,9 +18,10 @@ dtype = PrecisionManager.get_compute_dtype() """ -import tensorflow as tf -from typing import Optional import warnings +from typing import Optional + +import tensorflow as tf class PrecisionManager: diff --git a/src/quick_benchmark.py b/src/quick_benchmark.py index c52a9471b..4331ab7f4 100644 --- a/src/quick_benchmark.py +++ b/src/quick_benchmark.py @@ -14,16 +14,18 @@ python -m src.quick_benchmark --resolution 64 --batch_size 2 """ -import tensorflow as tf -import numpy as np -import time import argparse -from dataclasses import dataclass -from typing import Tuple, Optional +import os # Add src to path import sys -import os +import time +from dataclasses import dataclass +from typing import Optional + +import numpy as np +import tensorflow as tf + sys.path.insert(0, os.path.dirname(__file__)) from model_transforms import DeepCompressModel, DeepCompressModelV2, TransformConfig @@ -247,7 +249,7 @@ def run_benchmark( Returns: CompressionMetrics with results. """ - print(f"\nBenchmark Configuration:") + print("\nBenchmark Configuration:") print(f" Resolution: {resolution}x{resolution}x{resolution}") print(f" Batch size: {batch_size}") print(f" Model version: {model_version}") diff --git a/src/training_pipeline.py b/src/training_pipeline.py index 630285e82..e3d3709f4 100644 --- a/src/training_pipeline.py +++ b/src/training_pipeline.py @@ -1,9 +1,10 @@ -import tensorflow as tf -import numpy as np -import os import logging from pathlib import Path -from typing import Dict, Any, Optional +from typing import Dict + +import numpy as np +import tensorflow as tf + class TrainingPipeline: def __init__(self, config_path: str): @@ -104,44 +105,44 @@ def compute_focal_loss(self, y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor: def train(self, validate_every: int = 100): train_dataset = self.data_loader.load_training_data() val_dataset = self.data_loader.load_evaluation_data() - + step = 0 best_val_loss = float('inf') - + for epoch in range(self.config['training']['epochs']): self.logger.info(f"Epoch {epoch+1}/{self.config['training']['epochs']}") - + for batch in train_dataset: step += 1 losses = self._train_step(batch) - + with self.summary_writer.as_default(): for name, value in losses.items(): tf.summary.scalar(f'train/{name}', value, step=step) - + if step % validate_every == 0: val_losses = self._validate(val_dataset) - + with self.summary_writer.as_default(): for name, value in val_losses.items(): tf.summary.scalar(f'val/{name}', value, step=step) - + if val_losses['total_loss'] < best_val_loss: best_val_loss = val_losses['total_loss'] self.save_checkpoint('best_model') - + self.save_checkpoint(f'epoch_{epoch+1}') - + def _validate(self, val_dataset: tf.data.Dataset) -> Dict[str, float]: val_losses = [] for batch in val_dataset: losses = self._train_step(batch, training=False) val_losses.append({k: v.numpy() for k, v in losses.items()}) - + avg_losses = {} for metric in val_losses[0].keys(): avg_losses[metric] = float(tf.reduce_mean([x[metric] for x in val_losses])) - + return avg_losses def save_checkpoint(self, name: str): @@ -149,7 +150,7 @@ def save_checkpoint(self, name: str): checkpoint_path.mkdir(parents=True, exist_ok=True) self.model.save_weights(str(checkpoint_path / 'model.weights.h5')) self.entropy_model.save_weights(str(checkpoint_path / 'entropy.weights.h5')) - + for opt_name, optimizer in self.optimizers.items(): if optimizer.variables: opt_weights = [v.numpy() for v in optimizer.variables] @@ -165,7 +166,7 @@ def load_checkpoint(self, name: str): checkpoint_path = self.checkpoint_dir / name self.model.load_weights(str(checkpoint_path / 'model.weights.h5')) self.entropy_model.load_weights(str(checkpoint_path / 'entropy.weights.h5')) - + for opt_name, optimizer in self.optimizers.items(): opt_path = checkpoint_path / f'{opt_name}_optimizer.npy' if opt_path.exists() and optimizer.variables: @@ -181,16 +182,16 @@ def main(): parser.add_argument("config", type=str, help="Path to config file") parser.add_argument("--resume", type=str, help="Resume from checkpoint") args = parser.parse_args() - - logging.basicConfig(level=logging.INFO, + + logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') - + pipeline = TrainingPipeline(args.config) - + if args.resume: pipeline.load_checkpoint(args.resume) - + pipeline.train() if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/tests/conftest.py b/tests/conftest.py index aaff72ea3..a32418792 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,7 +1,9 @@ +from pathlib import Path + +import numpy as np import pytest import tensorflow as tf -import numpy as np -from pathlib import Path + def pytest_collection_modifyitems(items): @@ -59,7 +61,7 @@ def tf_config(): if gpus: for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) - + # Set random seeds for reproducibility tf.random.set_seed(42) - np.random.seed(42) \ No newline at end of file + np.random.seed(42) diff --git a/tests/test_attention_context.py b/tests/test_attention_context.py index a99953b46..a1e001bb6 100644 --- a/tests/test_attention_context.py +++ b/tests/test_attention_context.py @@ -1,18 +1,19 @@ """Tests for attention-based context model.""" -import tensorflow as tf -import pytest import sys from pathlib import Path +import pytest +import tensorflow as tf + # Add src to path for imports sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) from attention_context import ( - SparseAttention3D, - BidirectionalMaskTransformer, AttentionEntropyModel, - HybridAttentionEntropyModel + BidirectionalMaskTransformer, + HybridAttentionEntropyModel, + SparseAttention3D, ) diff --git a/tests/test_channel_context.py b/tests/test_channel_context.py index a1f859771..d27108939 100644 --- a/tests/test_channel_context.py +++ b/tests/test_channel_context.py @@ -1,13 +1,14 @@ """Tests for channel-wise context model.""" -import tensorflow as tf import sys from pathlib import Path +import tensorflow as tf + # Add src to path for imports sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) -from channel_context import SliceTransform, ChannelContext, ChannelContextEntropyModel +from channel_context import ChannelContext, ChannelContextEntropyModel, SliceTransform class TestSliceTransform(tf.test.TestCase): diff --git a/tests/test_colorbar.py b/tests/test_colorbar.py index c6b146786..048f50294 100644 --- a/tests/test_colorbar.py +++ b/tests/test_colorbar.py @@ -3,19 +3,22 @@ sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) -import pytest -import numpy as np -import matplotlib.pyplot as plt import json -from colorbar import get_colorbar, ColorbarConfig + +import matplotlib.pyplot as plt +import numpy as np +import pytest + +from colorbar import ColorbarConfig, get_colorbar + class TestColorbar: """Test suite for colorbar generation and color mapping functionality.""" - + def setup_method(self): self.vmin = 0 self.vmax = 100 - + def teardown_method(self): plt.close('all') @@ -23,7 +26,7 @@ def test_horizontal_colorbar(self): fig, cmap = get_colorbar(self.vmin, self.vmax, orientation='horizontal') assert len(fig.axes) > 0 assert callable(cmap) - + test_values = [0, 50, 100] colors = [cmap(val) for val in test_values] assert len(colors) == len(test_values) @@ -38,14 +41,14 @@ def test_vertical_colorbar(self): def test_custom_labels(self): labels = ['Low', 'Medium', 'High'] positions = [0, 50, 100] - + fig, cmap = get_colorbar( self.vmin, self.vmax, tick_labels=labels, tick_positions=positions ) - + cbar_ax = fig.axes[-1] tick_labels = [t.get_text() for t in cbar_ax.get_xticklabels()] assert tick_labels == labels @@ -59,7 +62,7 @@ def test_title_and_formatting(self): label_format='{:.2f}', tick_rotation=45 ) - + cbar_ax = fig.axes[-1] assert cbar_ax.get_xlabel() == title assert all(t.get_rotation() == 45 for t in cbar_ax.get_xticklabels()) @@ -70,11 +73,11 @@ def test_invalid_orientation(self): def test_color_mapping(self): fig, cmap = get_colorbar(self.vmin, self.vmax) - + color = cmap(50) assert len(color) == 4 assert all(0 <= c <= 1 for c in color) - + values = np.array([0, 50, 100]) colors = cmap(values) - assert colors.shape == (3, 4) \ No newline at end of file + assert colors.shape == (3, 4) diff --git a/tests/test_compress_octree.py b/tests/test_compress_octree.py index 61161773b..87ae819f7 100644 --- a/tests/test_compress_octree.py +++ b/tests/test_compress_octree.py @@ -1,13 +1,15 @@ import sys -import tensorflow as tf -import pytest -import numpy as np from pathlib import Path +import numpy as np +import pytest +import tensorflow as tf + sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) -from test_utils import create_mock_point_cloud, setup_test_environment from compress_octree import OctreeCompressor +from test_utils import create_mock_point_cloud, setup_test_environment + class TestOctreeCompressor(tf.test.TestCase): @pytest.fixture(autouse=True) diff --git a/tests/test_context_model.py b/tests/test_context_model.py index 9c3302c84..f78b8d7ea 100644 --- a/tests/test_context_model.py +++ b/tests/test_context_model.py @@ -1,14 +1,15 @@ """Tests for autoregressive context model.""" -import tensorflow as tf -import numpy as np import sys from pathlib import Path +import numpy as np +import tensorflow as tf + # Add src to path for imports sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) -from context_model import MaskedConv3D, AutoregressiveContext, ContextualEntropyModel +from context_model import AutoregressiveContext, ContextualEntropyModel, MaskedConv3D class TestMaskedConv3D(tf.test.TestCase): diff --git a/tests/test_data_loader.py b/tests/test_data_loader.py index feb12f6f1..539816675 100644 --- a/tests/test_data_loader.py +++ b/tests/test_data_loader.py @@ -1,13 +1,15 @@ import sys -import tensorflow as tf -import pytest -import numpy as np from pathlib import Path +import numpy as np +import pytest +import tensorflow as tf + sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) from data_loader import DataLoader + class TestDataLoader: @pytest.fixture def config(self): @@ -75,4 +77,4 @@ def test_batch_processing(self, data_loader, tmp_path, create_off_file, sample_p assert batch.shape[1:] == (resolution,) * 3 if __name__ == '__main__': - tf.test.main() \ No newline at end of file + tf.test.main() diff --git a/tests/test_ds_mesh_to_pc.py b/tests/test_ds_mesh_to_pc.py index 02faaae7a..bd5b7ffd3 100644 --- a/tests/test_ds_mesh_to_pc.py +++ b/tests/test_ds_mesh_to_pc.py @@ -1,20 +1,22 @@ +import os import sys import unittest -import numpy as np -import os from pathlib import Path +import numpy as np + sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) from ds_mesh_to_pc import ( + MeshData, + convert_mesh_to_point_cloud, + partition_point_cloud, read_off, sample_points_from_mesh, save_ply, - convert_mesh_to_point_cloud, - MeshData, - partition_point_cloud ) + class TestDsMeshToPc(unittest.TestCase): def setUp(self): self.test_off_file = "test.off" @@ -66,10 +68,10 @@ def test_read_off(self): def test_sample_points_from_mesh(self): mesh_data = MeshData(vertices=self.vertices, faces=None, vertex_normals=self.normals) points, normals = sample_points_from_mesh(mesh_data, num_points=3, compute_normals=True) - + self.assertEqual(points.shape, (3, 3)) self.assertEqual(normals.shape, (3, 3)) - + for point in points: self.assertTrue( np.any(np.all(np.abs(self.vertices - point) < 1e-5, axis=1)), @@ -93,7 +95,7 @@ def test_partition_point_cloud(self): block_size=0.3, min_points=2 ) - + self.assertGreater(len(blocks), 0) for block in blocks: self.assertIn('points', block) @@ -113,4 +115,4 @@ def test_end_to_end(self): self.assertTrue(os.path.exists(self.test_ply_file)) if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/tests/test_ds_pc_octree_blocks.py b/tests/test_ds_pc_octree_blocks.py index d1751ae86..ffc4917fe 100644 --- a/tests/test_ds_pc_octree_blocks.py +++ b/tests/test_ds_pc_octree_blocks.py @@ -1,12 +1,14 @@ import sys -import tensorflow as tf -import pytest from pathlib import Path +import pytest +import tensorflow as tf + sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) -from test_utils import create_mock_point_cloud, create_mock_ply_file, setup_test_environment from ds_pc_octree_blocks import PointCloudProcessor +from test_utils import create_mock_point_cloud, setup_test_environment + class TestPointCloudOctreeBlocks(tf.test.TestCase): @pytest.fixture(autouse=True) diff --git a/tests/test_entropy_model.py b/tests/test_entropy_model.py index c443c327d..f53e54fe9 100644 --- a/tests/test_entropy_model.py +++ b/tests/test_entropy_model.py @@ -4,7 +4,9 @@ sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) import tensorflow as tf -from entropy_model import PatchedGaussianConditional, EntropyModel + +from entropy_model import EntropyModel, PatchedGaussianConditional + class TestEntropyModel(tf.test.TestCase): diff --git a/tests/test_entropy_parameters.py b/tests/test_entropy_parameters.py index ab709aa62..d2de0bdd2 100644 --- a/tests/test_entropy_parameters.py +++ b/tests/test_entropy_parameters.py @@ -1,14 +1,15 @@ """Tests for entropy parameters network and mean-scale hyperprior.""" -import tensorflow as tf import sys from pathlib import Path +import tensorflow as tf + # Add src to path for imports sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) -from entropy_parameters import EntropyParameters, EntropyParametersWithContext from entropy_model import ConditionalGaussian, MeanScaleHyperprior +from entropy_parameters import EntropyParameters, EntropyParametersWithContext class TestEntropyParameters(tf.test.TestCase): diff --git a/tests/test_ev_run_render.py b/tests/test_ev_run_render.py index 54689fef4..afd49375f 100644 --- a/tests/test_ev_run_render.py +++ b/tests/test_ev_run_render.py @@ -1,23 +1,19 @@ -import sys +import json import os +import sys import tempfile + +import numpy as np import pytest -import yaml import tensorflow as tf -import numpy as np +import yaml from PIL import Image -import json # Add src directory to path sys.path.insert(0, str(os.path.join(os.path.dirname(__file__), '..', 'src'))) -from ev_run_render import ( - load_experiment_config, - RenderConfig, - CameraParams, - PointCloudRenderer, - save_rendered_image -) +from ev_run_render import CameraParams, PointCloudRenderer, RenderConfig, load_experiment_config, save_rendered_image + def create_test_config(): """Create test experiment configuration.""" @@ -57,24 +53,24 @@ def setup(self): self.config, self.dataset_dir, self.experiment_dir = create_test_config() self.render_config = RenderConfig(**self.config['render_config']) self.renderer = PointCloudRenderer(self.render_config) - + # Create test point cloud self.points = tf.random.uniform((100, 3), -1, 1, dtype=tf.float32) self.colors = tf.random.uniform((100, 3), 0, 1, dtype=tf.float32) - + yield - + # Cleanup os.rmdir(self.dataset_dir) os.rmdir(self.experiment_dir) - + def test_render_config(self): """Test render configuration.""" assert self.render_config.image_width == 128 assert self.render_config.image_height == 128 assert self.render_config.point_size == 2.0 assert self.render_config.color_map == 'plasma' - + def test_camera_params(self): """Test camera parameter handling.""" cam_config = self.config['data'][0]['camera_params'] @@ -84,12 +80,12 @@ def test_camera_params(self): up=np.array(cam_config['up']), fov=cam_config['fov'] ) - + assert np.allclose(camera.position, [0, 0, 5]) assert np.allclose(camera.target, [0, 0, 0]) assert np.allclose(camera.up, [0, 1, 0]) assert camera.fov == 45.0 - + def test_point_cloud_rendering(self): """Test basic point cloud rendering.""" camera = CameraParams( @@ -98,18 +94,18 @@ def test_point_cloud_rendering(self): up=np.array([0, 1, 0]), fov=45.0 ) - + image, render_info = self.renderer.render( self.points, colors=self.colors, camera=camera ) - + assert image.shape == (128, 128, 3) assert np.all(image >= 0) and np.all(image <= 1) assert 'camera' in render_info assert 'render_config' in render_info - + def test_save_rendered_image(self): """Test saving rendered image with metadata.""" with tempfile.TemporaryDirectory() as temp_dir: @@ -121,33 +117,33 @@ def test_save_rendered_image(self): fov=45.0 ) image, render_info = self.renderer.render(self.points, self.colors, camera) - + # Save image save_path = os.path.join(temp_dir, "test_render.png") bbox = (10, 10, 118, 118) save_rendered_image(image, render_info, save_path, bbox) - + # Check files exist assert os.path.exists(save_path) assert os.path.exists(save_path + ".meta.json") - + # Check image size img = Image.open(save_path) assert img.size == (108, 108) # Size after cropping - + # Check metadata with open(save_path + ".meta.json", 'r') as f: meta = json.load(f) assert 'camera' in meta assert 'render_config' in meta assert 'bbox' in meta - + def test_load_experiment_config(self): """Test experiment configuration loading.""" with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f: yaml.dump(self.config, f) config_path = f.name - + try: loaded_config = load_experiment_config(config_path) assert loaded_config['MPEG_DATASET_DIR'] == self.config['MPEG_DATASET_DIR'] @@ -155,17 +151,17 @@ def test_load_experiment_config(self): assert loaded_config['render_config'] == self.config['render_config'] finally: os.unlink(config_path) - + def test_missing_config_keys(self): """Test handling of missing configuration keys.""" invalid_config = { 'MPEG_DATASET_DIR': self.dataset_dir } - + with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f: yaml.dump(invalid_config, f) config_path = f.name - + try: with pytest.raises(ValueError): load_experiment_config(config_path) @@ -173,4 +169,4 @@ def test_missing_config_keys(self): os.unlink(config_path) if __name__ == '__main__': - pytest.main([__file__]) \ No newline at end of file + pytest.main([__file__]) diff --git a/tests/test_evaluation_pipeline.py b/tests/test_evaluation_pipeline.py index 652def1e3..a0f4bdae3 100644 --- a/tests/test_evaluation_pipeline.py +++ b/tests/test_evaluation_pipeline.py @@ -1,15 +1,17 @@ -import sys import json -import tensorflow as tf -import pytest -import numpy as np +import sys from pathlib import Path + +import numpy as np +import pytest +import tensorflow as tf import yaml sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) from evaluation_pipeline import EvaluationPipeline, EvaluationResult + class TestEvaluationPipeline: @pytest.fixture def config_path(self, tmp_path): diff --git a/tests/test_experiment.py b/tests/test_experiment.py index 01c80c73c..008c5af81 100644 --- a/tests/test_experiment.py +++ b/tests/test_experiment.py @@ -1,13 +1,15 @@ +import os import sys import unittest -import os -import yaml from pathlib import Path +import yaml + sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) from experiment import Experiment, ExperimentConfig + class TestExperimentConfig(unittest.TestCase): def setUp(self): self.test_config_path = "test_config.yml" @@ -26,8 +28,8 @@ def setUp(self): yaml.dump(self.test_config, f) def tearDown(self): - for path in [self.test_config_path, - self.test_config["dataset_path"], + for path in [self.test_config_path, + self.test_config["dataset_path"], self.test_config["experiment_dir"]]: if os.path.exists(path): if os.path.isdir(path): @@ -67,8 +69,8 @@ def setUp(self): self.experiment = Experiment(self.test_config_path) def tearDown(self): - for path in [self.test_config_path, - self.test_config["dataset_path"], + for path in [self.test_config_path, + self.test_config["dataset_path"], self.test_config["experiment_dir"]]: if os.path.exists(path): if os.path.isdir(path): @@ -81,4 +83,4 @@ def test_run_experiment(self): self.assertTrue(os.path.exists(self.test_config["experiment_dir"])) if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/tests/test_integration.py b/tests/test_integration.py index efef2e3b2..718e3a672 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -1,22 +1,18 @@ -import tensorflow as tf -import numpy as np -import pytest import sys from pathlib import Path +import pytest +import tensorflow as tf + # Add src to path for imports sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) -from test_utils import ( - setup_test_environment, - create_mock_point_cloud, - create_mock_voxel_grid, - create_test_dataset -) -from training_pipeline import TrainingPipeline -from evaluation_pipeline import EvaluationPipeline from data_loader import DataLoader +from evaluation_pipeline import EvaluationPipeline from model_transforms import DeepCompressModel, DeepCompressModelV2, TransformConfig +from test_utils import create_mock_point_cloud, create_mock_voxel_grid, create_test_dataset, setup_test_environment +from training_pipeline import TrainingPipeline + class TestIntegration(tf.test.TestCase): @pytest.fixture(autouse=True) @@ -327,4 +323,4 @@ def test_v2_gaussian_backward_compatible(self): if __name__ == '__main__': - tf.test.main() \ No newline at end of file + tf.test.main() diff --git a/tests/test_map_color.py b/tests/test_map_color.py index dbf9e8d60..045e140a1 100644 --- a/tests/test_map_color.py +++ b/tests/test_map_color.py @@ -1,12 +1,14 @@ +import os import sys import unittest -import numpy as np -import os from pathlib import Path +import numpy as np + sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) -from map_color import load_point_cloud, load_colors, transfer_colors, save_colored_point_cloud +from map_color import load_colors, load_point_cloud, save_colored_point_cloud, transfer_colors + class TestMapColor(unittest.TestCase): """Test suite for point cloud color mapping operations.""" @@ -92,4 +94,4 @@ def test_save_colored_point_cloud(self): self.assertIn("property uchar blue", content) if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/tests/test_model_transforms.py b/tests/test_model_transforms.py index 740babb42..03c8c2bf5 100644 --- a/tests/test_model_transforms.py +++ b/tests/test_model_transforms.py @@ -1,21 +1,23 @@ -import tensorflow as tf -import pytest import sys from pathlib import Path +import pytest +import tensorflow as tf + # Add src to path for imports sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) -from test_utils import create_mock_voxel_grid from model_transforms import ( CENICGDN, - SpatialSeparableConv, AnalysisTransform, - SynthesisTransform, DeepCompressModel, DeepCompressModelV2, - TransformConfig + SpatialSeparableConv, + SynthesisTransform, + TransformConfig, ) +from test_utils import create_mock_voxel_grid + class TestModelTransforms(tf.test.TestCase): @pytest.fixture(autouse=True) @@ -52,9 +54,8 @@ def test_spatial_separable_conv(self): input_tensor = tf.random.uniform((2, 32, 32, 32, 32)) output = conv(input_tensor) self.assertEqual(output.shape[-1], 64) - + standard_params = 27 * 32 * 64 - separable_params = (3 * 32 * 32 + 9 * 32 * 64) self.assertLess(len(conv.trainable_variables[0].numpy().flatten()), standard_params) def test_analysis_transform(self): @@ -115,8 +116,8 @@ def test_model_save_load(self): input_tensor = create_mock_voxel_grid(self.resolution, self.batch_size) x_hat1, y1, y_hat1, z1 = model(input_tensor, training=False) - import tempfile import os + import tempfile with tempfile.TemporaryDirectory() as tmp_dir: # Keras 3 requires .weights.h5 extension for save_weights save_path = os.path.join(tmp_dir, 'model.weights.h5') @@ -317,4 +318,4 @@ def test_different_entropy_models_produce_different_rates(self): if __name__ == "__main__": - tf.test.main() \ No newline at end of file + tf.test.main() diff --git a/tests/test_mp_report.py b/tests/test_mp_report.py index 50aa6356e..a7e50bfc3 100644 --- a/tests/test_mp_report.py +++ b/tests/test_mp_report.py @@ -1,10 +1,11 @@ -import sys -import pytest -import os import json +import os +import sys from pathlib import Path from tempfile import TemporaryDirectory +import pytest + sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) from mp_report import ExperimentReporter, load_experiment_results diff --git a/tests/test_octree_coding.py b/tests/test_octree_coding.py index e99a73568..9a876e72e 100644 --- a/tests/test_octree_coding.py +++ b/tests/test_octree_coding.py @@ -5,9 +5,12 @@ sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) import unittest + import tensorflow as tf + from octree_coding import OctreeCoder, OctreeConfig + class TestOctreeCoder(unittest.TestCase): def setUp(self): diff --git a/tests/test_parallel_process.py b/tests/test_parallel_process.py index 8a97cbd8f..c57274da4 100644 --- a/tests/test_parallel_process.py +++ b/tests/test_parallel_process.py @@ -1,19 +1,15 @@ +import subprocess import sys +import time +import unittest +from concurrent.futures import TimeoutError from pathlib import Path +from unittest.mock import MagicMock, patch sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) -import unittest -import time -import subprocess -from unittest.mock import patch, MagicMock -from concurrent.futures import TimeoutError -from parallel_process import ( - parallel_process, - Popen, - ProcessTimeoutError, - ProcessResult -) +from parallel_process import Popen, ProcessResult, ProcessTimeoutError, parallel_process + def square(x): """Simple square function for testing.""" @@ -42,17 +38,17 @@ def test_parallel_process_basic(self): def test_parallel_process_timeout(self): """Test timeout functionality.""" params = [0.1, 0.1, 2.0] # Last task will timeout - + with self.assertRaises(TimeoutError): - parallel_process(slow_function, params, + parallel_process(slow_function, params, num_parallel=2, timeout=1.0) def test_parallel_process_retries(self): """Test retry functionality for failing tasks.""" params = [1, 2, 3] # 2 will fail - + with self.assertRaises(ValueError): - parallel_process(failing_function, params, + parallel_process(failing_function, params, num_parallel=2, max_retries=2) @patch("subprocess.Popen") @@ -78,7 +74,7 @@ def test_popen_cleanup(self, mock_popen): mock_popen.return_value = mock_process cmd = ["echo", "test"] - with Popen(cmd) as process: + with Popen(cmd) as _: pass # Context manager should handle cleanup mock_process.terminate.assert_called_once() @@ -93,11 +89,11 @@ def test_process_result_dataclass(self): success=True, error=None ) - + self.assertEqual(result.index, 0) self.assertEqual(result.result, 42) self.assertTrue(result.success) self.assertIsNone(result.error) if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/tests/test_performance.py b/tests/test_performance.py index 315b28882..b7a114079 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -9,15 +9,16 @@ Run with: pytest tests/test_performance.py -v """ -import sys import os +import sys + sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src')) -import pytest -import tensorflow as tf -import numpy as np import time +import numpy as np +import pytest +import tensorflow as tf # ============================================================================= # Fixtures diff --git a/tests/test_point_cloud_metricss.py b/tests/test_point_cloud_metricss.py index 54cfc25f9..cb6bcac58 100644 --- a/tests/test_point_cloud_metricss.py +++ b/tests/test_point_cloud_metricss.py @@ -1,14 +1,16 @@ import sys -import tensorflow as tf -import pytest -import numpy as np from pathlib import Path +import numpy as np +import pytest +import tensorflow as tf + sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) -from point_cloud_metrics import calculate_metrics, calculate_chamfer_distance, calculate_d1_metric +from point_cloud_metrics import calculate_chamfer_distance, calculate_d1_metric, calculate_metrics from test_utils import create_mock_point_cloud + class TestPointCloudMetrics(tf.test.TestCase): @pytest.fixture(autouse=True) def setup(self): diff --git a/tests/test_training_pipeline.py b/tests/test_training_pipeline.py index 2d75272fb..4a5b4290a 100644 --- a/tests/test_training_pipeline.py +++ b/tests/test_training_pipeline.py @@ -1,14 +1,15 @@ import sys -import tensorflow as tf -import pytest -import numpy as np from pathlib import Path + +import pytest +import tensorflow as tf import yaml sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) from training_pipeline import TrainingPipeline + class TestTrainingPipeline: @pytest.fixture def config_path(self, tmp_path): @@ -40,7 +41,7 @@ def config_path(self, tmp_path): 'checkpoint_dir': str(tmp_path / 'checkpoints') } } - + config_file = tmp_path / 'config.yml' with open(config_file, 'w') as f: yaml.dump(config, f) @@ -60,10 +61,10 @@ def test_initialization(self, pipeline): def test_compute_focal_loss(self, pipeline): batch_size = 4 resolution = 32 - + y_true = tf.cast(tf.random.uniform((batch_size, resolution, resolution, resolution)) > 0.5, tf.float32) y_pred = tf.random.uniform((batch_size, resolution, resolution, resolution)) - + loss = pipeline.compute_focal_loss(y_true, y_pred) assert loss.shape == () assert loss >= 0 @@ -74,13 +75,13 @@ def test_train_step(self, pipeline, training): batch_size = 1 resolution = 16 point_cloud = tf.cast(tf.random.uniform((batch_size, resolution, resolution, resolution)) > 0.5, tf.float32) - + losses = pipeline._train_step(point_cloud, training=training) - + assert 'focal_loss' in losses assert 'entropy_loss' in losses assert 'total_loss' in losses - + for loss_name, loss_value in losses.items(): assert not tf.math.is_nan(loss_value) assert loss_value >= 0 @@ -94,18 +95,18 @@ def test_save_load_checkpoint(self, pipeline, tmp_path): checkpoint_name = 'test_checkpoint' pipeline.save_checkpoint(checkpoint_name) - + checkpoint_dir = Path(pipeline.checkpoint_dir) / checkpoint_name assert (checkpoint_dir / 'model.weights.h5').exists() assert (checkpoint_dir / 'entropy.weights.h5').exists() - + new_pipeline = TrainingPipeline(pipeline.config_path) # Build the new model before loading weights new_pipeline.model(dummy, training=False) y2 = new_pipeline.model.analysis(dummy) new_pipeline.entropy_model(y2, training=False) new_pipeline.load_checkpoint(checkpoint_name) - + for w1, w2 in zip(pipeline.model.weights, new_pipeline.model.weights): tf.debugging.assert_equal(w1, w2) @@ -116,15 +117,15 @@ def test_training_loop(self, pipeline, tmp_path): def create_sample_batch(): return tf.cast(tf.random.uniform((batch_size, resolution, resolution, resolution)) > 0.5, tf.float32) - + dataset = tf.data.Dataset.from_tensors(create_sample_batch()).repeat(3) - + pipeline.data_loader.load_training_data = lambda: dataset pipeline.data_loader.load_evaluation_data = lambda: dataset - + pipeline.config['training']['epochs'] = 2 pipeline.train(validate_every=2) - + checkpoint_dir = Path(pipeline.checkpoint_dir) assert len(list(checkpoint_dir.glob('epoch_*'))) > 0 - assert (checkpoint_dir / 'best_model').exists() \ No newline at end of file + assert (checkpoint_dir / 'best_model').exists() diff --git a/tests/test_utils.py b/tests/test_utils.py index dd2a19e96..48e228cc4 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,8 +1,9 @@ -import tensorflow as tf -import numpy as np from pathlib import Path +from typing import Any, Dict, Optional + +import tensorflow as tf import yaml -from typing import Optional, Dict, Any, List, Tuple + def create_mock_point_cloud(num_points: int = 1000) -> tf.Tensor: """Create a mock point cloud for testing.""" @@ -21,14 +22,14 @@ def create_mock_normals(points: tf.Tensor) -> tf.Tensor: normals = tf.random.normal(points.shape) return tf.nn.l2_normalize(normals, axis=1) -def create_mock_ply_file(filepath: Path, points: Optional[tf.Tensor] = None, +def create_mock_ply_file(filepath: Path, points: Optional[tf.Tensor] = None, normals: Optional[tf.Tensor] = None): """Create a mock PLY file with given points and normals.""" if points is None: points = create_mock_point_cloud() if normals is None: normals = create_mock_normals(points) - + with open(filepath, 'w') as f: f.write("ply\n") f.write("format ascii 1.0\n") @@ -36,14 +37,14 @@ def create_mock_ply_file(filepath: Path, points: Optional[tf.Tensor] = None, f.write("property float x\n") f.write("property float y\n") f.write("property float z\n") - + if normals is not None: f.write("property float nx\n") f.write("property float ny\n") f.write("property float nz\n") - + f.write("end_header\n") - + for i in range(len(points)): line = f"{points[i, 0]} {points[i, 1]} {points[i, 2]}" if normals is not None: @@ -63,10 +64,10 @@ def create_test_off_file(filepath: Path, mesh: Optional[Dict[str, tf.Tensor]] = """Create a test OFF file with mesh data.""" if mesh is None: mesh = create_test_mesh() - + vertices = mesh['vertices'] faces = mesh['faces'] - + with open(filepath, 'w') as f: f.write("OFF\n") f.write(f"{len(vertices)} {len(faces)} 0\n") @@ -77,7 +78,7 @@ def create_test_off_file(filepath: Path, mesh: Optional[Dict[str, tf.Tensor]] = for face in faces: f.write(f"3 {face[0]} {face[1]} {face[2]}\n") -def create_test_dataset(batch_size: int, resolution: int, +def create_test_dataset(batch_size: int, resolution: int, num_batches: int = 10) -> tf.data.Dataset: """Create a test dataset with proper shape.""" return tf.data.Dataset.from_tensor_slices( @@ -123,28 +124,28 @@ def create_test_config(tmp_path: Path) -> Dict[str, Any]: def setup_test_environment(tmp_path: Path) -> Dict[str, Any]: """Set up a complete test environment with files and configs.""" config = create_test_config(tmp_path) - + # Create directories for key in ['modelnet40_path', 'ivfb_path']: Path(config['data'][key]).mkdir(parents=True, exist_ok=True) - + # Create test files test_files = { 'mesh': Path(config['data']['modelnet40_path']) / "test.off", 'point_cloud': Path(config['data']['ivfb_path']) / "test.ply", 'blocks': Path(config['evaluation']['output_dir']) / "blocks" } - + mesh = create_test_mesh() points = create_mock_point_cloud() - + create_test_off_file(test_files['mesh'], mesh) create_mock_ply_file(test_files['point_cloud'], points) - + config_path = tmp_path / 'config.yml' with open(config_path, 'w') as f: yaml.dump(config, f) - + return { 'config': config, 'config_path': str(config_path), @@ -158,10 +159,10 @@ def __init__(self): super().__init__() self.batch_end_called = 0 self.epoch_end_called = 0 - + def on_batch_end(self, batch, logs=None): self.batch_end_called += 1 - + def on_epoch_end(self, epoch, logs=None): self.epoch_end_called += 1 @@ -175,4 +176,4 @@ def compute_mock_metrics(predicted: tf.Tensor, target: tf.Tensor) -> Dict[str, t } if __name__ == "__main__": - tf.test.main() \ No newline at end of file + tf.test.main()