feat: add version 2 - complex field support

.gitignore

@@ -1,9 +1,11 @@
 .venv
 __pycache__
+**/*.pth
 checkpoints/
 logs/
 .DS_Store
 .env
 .pytest_cache/
 .ruff_cache/
-.cursor/
+.cursor/
+**/*egg-info*

README.md

@@ -1,10 +1,34 @@
-# HOLO-PASWIN: In-Line Holographical Physics-Aware SWIN Transformer
+# HoloPASWIN v2
 
-A research-grade deep learning project for eliminating the **twin-image problem in in-line holography** using a **physics-aware Swin-UNet architecture** trained with synthetic holograms generated via the Angular Spectrum Method.
+Physics-Aware Swin Transformer for eliminating twin-image artifacts in in-line holography.
 
-Build a modern, high-performing, physics-aware neural architecture capable of reconstructing complex object fields ( $\hat{O} = \hat{O}_{real} + i\hat{O}_{imag}$ ) from inline holograms while robustly suppressing **twin-image artifacts**, **noise**, and **aberrations**.
+## Installation
 
-## Development Setup
+This project uses [uv](https://github.com/astral-sh/uv) for dependency management.
+
+1.  **Install uv** (if not already installed):
+    ```bash
+    curl -LsSf https://astral.sh/uv/install.sh | sh
+    ```
+
+2.  **Sync Dependencies**:
+    Navigate to the `holopaswin` directory and run:
+    ```bash
+    uv sync
+    ```
+    This will create a virtual environment and install all locked dependencies from `uv.lock`.
+
+3.  **Training**:
+    To start training, run:
+    ```bash
+    uv run src/train.py
+    ```
+
+## Development
+
+-   **HOLO-PASWIN v2** builds upon Swin Transformer U-Net architecture.
+-   It accepts Hologram Intensity and outputs Complex Object (Phase & Amplitude).
+-   Dataset is loaded from efficient Parquet files.
 
 ### Installing Git Hooks
 

pyproject.toml

@@ -38,6 +38,8 @@ dependencies = [
     "scikit-image>=0.20.0",
     "matplotlib>=3.7.0",
     "tqdm>=4.65.0",
+    "pandas>=2.3.3",
+    "pyarrow>=22.0.0",
 ]
 
 [project.optional-dependencies]
@@ -110,3 +112,10 @@ module = [
 ]
 ignore_missing_imports = true
 
+[dependency-groups]
+dev = [
+    "mypy>=1.19.0",
+    "ruff>=0.14.9",
+    "types-tqdm>=4.67.0.20250809",
+]
+

results/experiment3/detailed_test.py

@@ -0,0 +1,172 @@
+"""Detailed evaluation script for Phase and Magnitude analysis.
+
+This script performs a deep-dive analysis on a small subset of the test dataset (10 samples).
+It computes and reports SSIM and PSNR separately for the Magnitude (Amplitude) and Phase
+domains of the reconstructed complex field. It also generates a comparative visualization.
+"""
+
+import os
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.utils.data
+from skimage.metrics import peak_signal_noise_ratio as psnr
+from skimage.metrics import structural_similarity as ssim
+
+from holopaswin.dataset import HoloDataset
+from holopaswin.model import HoloPASWIN
+
+# Configuration
+IMG_SIZE = 224
+WAVELENGTH = 532e-9
+PIXEL_SIZE = 4.65e-6
+Z_DIST = 0.02
+BATCH_SIZE = 1
+
+MODEL_PATH = "results/experiment3/holopaswin_exp3.pth"
+DATA_DIR = "../hologen/inline-digital-holography-v2"
+OUTPUT_IMG = "detailed_test_comparison.png"
+NUM_SAMPLES = 10
+
+
+def evaluate_detailed_metrics(model_path: str, data_dir: str) -> None:  # noqa: PLR0915
+    """Run detailed evaluation on phase and magnitude."""
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if torch.backends.mps.is_available():
+        device = torch.device("mps")
+    print(f"--- Starting Detailed Evaluation on {device} ---")
+
+    # 1. Load Model
+    if not os.path.exists(model_path):
+        print(f"Model not found at {model_path}")
+        return
+
+    model = HoloPASWIN(IMG_SIZE, WAVELENGTH, PIXEL_SIZE, Z_DIST).to(device)
+    try:
+        model.load_state_dict(torch.load(model_path, map_location=device))
+    except Exception as e:
+        print(f"Error loading weights: {e}")
+        return
+    model.eval()
+
+    # 2. Load Dataset
+    print(f"Loading dataset from {data_dir}...")
+    try:
+        dataset = HoloDataset(data_dir, target_size=IMG_SIZE)
+    except Exception as e:
+        print(f"Error loading dataset: {e}")
+        return
+
+    if len(dataset) == 0:
+        print("Dataset is empty!")
+        return
+
+    # Select random samples
+    indices = np.random.choice(len(dataset), NUM_SAMPLES, replace=False)
+    print(f"Selected indices: {indices}")
+
+    # Prepare plotting
+    fig, axes = plt.subplots(NUM_SAMPLES, 5, figsize=(20, 4 * NUM_SAMPLES))
+    cols = ["Hologram", "Pred Amp", "GT Amp", "Pred Phase", "GT Phase"]
+    for ax, col in zip(axes[0], cols, strict=False):
+        ax.set_title(col, fontsize=12, fontweight="bold")
+
+    print("\n" + "=" * 80)
+    print(f"{'Sample':<8} | {'Amp SSIM':<10} | {'Amp PSNR':<10} | {'Phase SSIM':<10} | {'Phase PSNR':<10}")
+    print("-" * 80)
+
+    avg_amp_ssim = 0.0
+    avg_amp_psnr = 0.0
+    avg_phase_ssim = 0.0
+    avg_phase_psnr = 0.0
+
+    with torch.no_grad():
+        for i, idx in enumerate(indices):
+            # Load Data
+            holo_t, gt_obj_t = dataset[idx]  # holo: (1,H,W), gt: (2,H,W)
+            holo_in = holo_t.unsqueeze(0).to(device)
+
+            # Inference
+            clean_pred, _ = model(holo_in)
+
+            # --- PROCESS OUTPUTS ---
+            # 1. Prediction Complex Field
+            pred_c = torch.complex(clean_pred[:, 0, :, :], clean_pred[:, 1, :, :])
+            pred_amp = torch.abs(pred_c).squeeze().cpu().numpy()
+            pred_phase = torch.angle(pred_c).squeeze().cpu().numpy()
+
+            # 2. Ground Truth Complex Field
+            gt_c = torch.complex(gt_obj_t[0], gt_obj_t[1])
+            gt_amp = torch.abs(gt_c).numpy()
+            gt_phase = torch.angle(gt_c).numpy()
+
+            # 3. Hologram
+            img_holo = holo_t.squeeze().numpy()
+
+            # --- METRICS CALCULATION ---
+
+            # Amplitude Metrics
+            # Range: [0, 1.2] roughly. We assume physics range 1.0 for standard metric scaling,
+            # or we can use dynamic range. Using 1.0 is standard if data is mostly [0,1].
+            # However, phase is [-pi, pi], range ~ 6.28.
+
+            data_range_amp = 1.2
+            data_range_phase = 2 * np.pi
+
+            score_amp_ssim = ssim(gt_amp, pred_amp, data_range=data_range_amp)  # type: ignore[no-untyped-call]
+            score_amp_psnr = psnr(gt_amp, pred_amp, data_range=data_range_amp)  # type: ignore[no-untyped-call]
+
+            score_phase_ssim = ssim(gt_phase, pred_phase, data_range=data_range_phase)  # type: ignore[no-untyped-call]
+            score_phase_psnr = psnr(gt_phase, pred_phase, data_range=data_range_phase)  # type: ignore[no-untyped-call]
+
+            # Accumulate
+            avg_amp_ssim += score_amp_ssim
+            avg_amp_psnr += score_amp_psnr
+            avg_phase_ssim += score_phase_ssim
+            avg_phase_psnr += score_phase_psnr
+
+            print(
+                f"{i:<8} | {score_amp_ssim:.4f}     | {score_amp_psnr:.2f} dB    | {score_phase_ssim:.4f}       | {score_phase_psnr:.2f} dB"
+            )
+
+            # --- PLOTTING ---
+            ax_row = axes[i]
+
+            # Hologram
+            ax_row[0].imshow(img_holo, cmap="gray")
+            ax_row[0].axis("off")
+
+            # Amplitude
+            ax_row[1].imshow(pred_amp, cmap="inferno", vmin=0, vmax=1.2)
+            ax_row[1].axis("off")
+            ax_row[2].imshow(gt_amp, cmap="inferno", vmin=0, vmax=1.2)
+            ax_row[2].axis("off")
+
+            # Phase
+            ax_row[3].imshow(pred_phase, cmap="twilight", vmin=-np.pi, vmax=np.pi)
+            ax_row[3].axis("off")
+            ax_row[4].imshow(gt_phase, cmap="twilight", vmin=-np.pi, vmax=np.pi)
+            ax_row[4].axis("off")
+
+    # Averages
+    avg_amp_ssim /= NUM_SAMPLES
+    avg_amp_psnr /= NUM_SAMPLES
+    avg_phase_ssim /= NUM_SAMPLES
+    avg_phase_psnr /= NUM_SAMPLES
+
+    print("-" * 80)
+    print(
+        f"{'AVERAGE':<8} | {avg_amp_ssim:.4f}     | {avg_amp_psnr:.2f} dB    | {avg_phase_ssim:.4f}       | {avg_phase_psnr:.2f} dB"
+    )
+    print("=" * 80)
+
+    plt.tight_layout()
+    plt.savefig(OUTPUT_IMG, dpi=150)
+    print(f"\nSaved comparison figure to {OUTPUT_IMG}")
+    plt.show()
+
+
+if __name__ == "__main__":
+    evaluate_detailed_metrics(MODEL_PATH, DATA_DIR)

results/experiment3/detailed_test.txt

@@ -0,0 +1,25 @@
+# uv run src/detailed_test.py
+
+--- Starting Detailed Evaluation on mps ---
+Loading dataset from ../hologen/inline-digital-holography-v2...
+Found 48 hologram parquet files.
+Indexing dataset...
+Indexed 12000 samples.
+Selected indices: [ 2007  4420  3664  1868  6034  3392  9551  9368  5093 11043]
+
+================================================================================
+Sample   | Amp SSIM   | Amp PSNR   | Phase SSIM | Phase PSNR
+--------------------------------------------------------------------------------
+0        | 0.6985     | 1.58 dB    | 0.6088       | 41.12 dB
+1        | 0.7483     | 1.58 dB    | 0.6113       | 41.43 dB
+2        | 0.7408     | 1.58 dB    | 0.6102       | 41.31 dB
+3        | 0.7042     | 1.58 dB    | 0.6103       | 41.15 dB
+4        | 0.7441     | 1.58 dB    | 0.6096       | 41.25 dB
+5        | 0.7667     | 1.58 dB    | 0.6129       | 41.80 dB
+6        | 0.6905     | 1.58 dB    | 0.6097       | 40.96 dB
+7        | 0.6918     | 1.58 dB    | 0.6094       | 41.09 dB
+8        | 0.6832     | 1.58 dB    | 0.6093       | 40.88 dB
+9        | 0.6937     | 1.49 dB    | 0.6079       | 41.12 dB
+--------------------------------------------------------------------------------
+AVERAGE  | 0.7162     | 1.57 dB    | 0.6099       | 41.21 dB
+================================================================================

results/experiment5/test_results.txt

@@ -0,0 +1,34 @@
+--- Starting Evaluation on mps ---
+Loading dataset from results/experiment5/test-dataset...
+Found 2 hologram parquet files.
+Indexing dataset...
+Indexed 496 samples.
+Evaluating on 496 random samples...
+Sample 50/496 -> Amp SSIM: 0.9440 | Phase SSIM: 0.9868
+Sample 100/496 -> Amp SSIM: 0.9860 | Phase SSIM: 0.9882
+Sample 150/496 -> Amp SSIM: 0.9520 | Phase SSIM: 0.9902
+Sample 200/496 -> Amp SSIM: 0.9453 | Phase SSIM: 0.9884
+Sample 250/496 -> Amp SSIM: 0.9786 | Phase SSIM: 0.9860
+Sample 300/496 -> Amp SSIM: 0.9379 | Phase SSIM: 0.9832
+Sample 350/496 -> Amp SSIM: 0.9851 | Phase SSIM: 0.9883
+Sample 400/496 -> Amp SSIM: 0.9835 | Phase SSIM: 0.9898
+Sample 450/496 -> Amp SSIM: 0.9444 | Phase SSIM: 0.9819
+
+============================================================
+   FINAL MODEL ACCURACY REPORT   
+============================================================
+Evaluated on: 496 samples
+------------------------------------------------------------
+AMPLITUDE DOMAIN (Absorption/Objects)
+  MSE:  0.000085 (±0.000043)
+  SSIM: 0.9642   (±0.0190)
+  PSNR: 42.88 dB  (±2.37)
+------------------------------------------------------------
+PHASE DOMAIN (Thickness/Refractive Index)
+  MSE:  0.000608 (±0.000229)
+  SSIM: 0.9868   (±0.0026)
+  PSNR: 48.44 dB  (±1.68)
+------------------------------------------------------------
+COMPLEX DOMAIN (Overall Fidelity)
+  MSE:  0.000659 (±0.000224)
+============================================================