Merge branch 'main' of github.com:tristan-deep/semantic-diffusion-echo-dehazing

.gitignore

@@ -5,3 +5,4 @@ temp/
 *.pdf
 *.hash
 *.npz
+sweep_results/

eval.py

@@ -208,7 +208,7 @@ def calculate_final_score(aggregates):
         return 0
 
 
-def main(folder: str, noisy_folder: str, roi_folder: str, reference_folder: str):
+def evaluate(folder: str, noisy_folder: str, roi_folder: str, reference_folder: str):
     """Evaluate the dehazing algorithm.
 
     Args:
@@ -294,6 +294,25 @@ def main(folder: str, noisy_folder: str, roi_folder: str, reference_folder: str)
     fid_score = calculate_fid_score(fid_image_paths, str(reference_folder))
     print(f"FID between {folder} and {reference_folder}: {fid_score:.3f}")
 
+    # Create aggregates dictionary for final score calculation
+    aggregates = {
+        "fid": float(fid_score),
+        "cnr_mean": float(np.mean(metrics["CNR"])),
+        "cnr_std": float(np.std(metrics["CNR"])),
+        "gcnr_mean": float(np.mean(metrics["gCNR"])),
+        "gcnr_std": float(np.std(metrics["gCNR"])),
+        "ks_roi1_ksstatistic_mean": float(np.mean(metrics["KS_A"])),
+        "ks_roi1_ksstatistic_std": float(np.std(metrics["KS_A"])),
+        "ks_roi2_ksstatistic_mean": float(np.mean(metrics["KS_B"])),
+        "ks_roi2_ksstatistic_std": float(np.std(metrics["KS_B"])),
+    }
+
+    # Calculate final score
+    final_score = calculate_final_score(aggregates)
+    aggregates["final_score"] = float(final_score)
+
+    return aggregates
+
 
 if __name__ == "__main__":
-    tyro.cli(main)
+    tyro.cli(evaluate)

plots.py

@@ -1,5 +1,10 @@
+import json
+from pathlib import Path
+from typing import Any, Dict, List
+
 import matplotlib.pyplot as plt
 import numpy as np
+import tyro
 from keras import ops
 from matplotlib.patches import PathPatch
 from matplotlib.path import Path as pltPath
@@ -252,3 +257,102 @@ def plot_metrics(metrics, limits, out_path):
         bbox_to_anchor=(0.5, 1.02),
     )
     return fig
+
+
+def plot_optimization_history_from_json(
+    trials_data: List[Dict[str, Any]], output_path: Path, method: str
+):
+    """Plot optimization history directly from JSON data."""
+
+    # Extract completed trials with values
+    completed_trials = [
+        t for t in trials_data if t["state"] == "COMPLETE" and t["value"] is not None
+    ]
+
+    if not completed_trials:
+        print("No completed trials found!")
+        return
+
+    # Sort by trial number
+    completed_trials.sort(key=lambda x: x["number"])
+
+    trial_numbers = [t["number"] for t in completed_trials]
+    values = [t["value"] for t in completed_trials]
+
+    # Apply seaborn styling
+    plt.style.use("seaborn-v0_8-darkgrid")
+
+    # Create the plot
+    fig, ax = plt.subplots(figsize=(5, 3), dpi=600)
+
+    # Plot all trial values with styling similar to plots.py
+    ax.scatter(
+        trial_numbers,
+        values,
+        c="#0057b7",
+        alpha=0.6,
+        s=30,
+        edgecolor="black",
+        linewidth=0.5,
+    )
+
+    # Plot best value so far
+    best_values = []
+    current_best = values[0]
+    for val in values:
+        if val > current_best:  # Assuming maximization
+            current_best = val
+        best_values.append(current_best)
+
+    ax.plot(
+        trial_numbers,
+        best_values,
+        color="#d62d20",
+        linewidth=2.5,
+        label="Best Value",
+        marker="o",
+        markersize=4,
+        markevery=max(1, len(trial_numbers) // 20),
+    )
+
+    ax.set_xlabel("Trial", fontsize=11)
+    ax.set_ylabel("Objective Value", fontsize=11)
+    # ax.set_title("Optimization History", fontsize=12)
+    ax.legend(fontsize=10, frameon=False)
+
+    # Remove top and right spines like in plots.py
+    ax.spines["top"].set_visible(False)
+    ax.spines["right"].set_visible(False)
+    ax.tick_params(axis="both", which="major", labelsize=9)
+
+    # Save plot
+    fig.savefig(
+        output_path / f"optimization_history_{method}.png", dpi=600, bbox_inches="tight"
+    )
+    fig.savefig(
+        output_path / f"optimization_history_{method}.pdf", dpi=600, bbox_inches="tight"
+    )
+    plt.close(fig)
+
+
+def main(json_file: str, output_dir: str = "plots", method: str = "semantic_dps"):
+    json_path = Path(json_file)
+    if not json_path.exists():
+        raise FileNotFoundError(f"JSON file not found: {json_file}")
+
+    # Load trial data
+    with open(json_path, "r") as f:
+        trials_data = json.load(f)
+
+    print(f"Loaded {len(trials_data)} trials from {json_file}")
+
+    # Create output directory
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    print("Generating optimization history plot...")
+    plot_optimization_history_from_json(trials_data, output_path, method)
+
+
+if __name__ == "__main__":
+    tyro.cli(main)

sweeper.py

@@ -0,0 +1,418 @@
+"""
+
+NOTE: pip install optuna
+
+"""
+
+import dataclasses
+import json
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import jax
+import numpy as np
+import optuna
+import tyro
+import yaml
+import zea
+from keras import ops
+from PIL import Image
+from zea import init_device, log
+
+from eval import evaluate
+from main import init, run
+
+
+def load_images_from_dir(input_folder):
+    """Load images from directory, similar to main.py implementation."""
+    paths = list(Path(input_folder).glob("*.png"))
+
+    images = []
+    for path in paths:
+        image = zea.io_lib.load_image(path)
+        images.append(image)
+
+    if len(images) == 0:
+        raise ValueError(f"No PNG images found in {input_folder}")
+
+    images = ops.stack(images, axis=0)
+    return images, paths
+
+
+def save_images_to_temp_dir(images, image_paths, prefix=""):
+    """Save numpy arrays as PNG images to a temporary directory."""
+    temp_dir = tempfile.mkdtemp(prefix=prefix)
+    temp_dir_path = Path(temp_dir)
+
+    for i, (img, path) in enumerate(zip(images, image_paths)):
+        # Get the filename from the original path
+        filename = Path(path).name
+
+        # Convert image to uint8 if needed
+        if img.dtype != np.uint8:
+            # Assume image is in [0, 1] range and convert to [0, 255]
+            if img.max() <= 1.0:
+                img = (img * 255).astype(np.uint8)
+            else:
+                img = img.astype(np.uint8)
+
+        # Ensure image is 2D or 3D
+        if len(img.shape) == 3 and img.shape[-1] == 1:
+            img = img.squeeze(-1)
+
+        # Save as PNG
+        img_pil = Image.fromarray(img)
+        save_path = temp_dir_path / filename
+        img_pil.save(save_path)
+
+    return str(temp_dir_path)
+
+
+@dataclasses.dataclass
+class SweeperConfig:
+    """Configuration for hyperparameter sweeping with Optuna."""
+
+    # Required paths - no defaults
+    input_image_dir: str  # Path to input hazy images
+    roi_folder: str  # Path to ROI mask images
+    reference_folder: str  # Path to reference/ground truth images
+    base_config_path: str = "configs/semantic_dps.yaml"
+
+    # Base configuration
+    method: str = "semantic_dps"  # Which method to optimize
+    broad_sweep: bool = False  # Choose between broad or narrow sweep
+
+    # Optuna settings
+    study_name: str = "dehaze_optimization"
+    storage: Optional[str] = None  # e.g., "sqlite:///dehaze_study.db" for persistence
+    n_trials: int = 100
+
+    # Optimization settings
+    objective_metric: str = "final_score"  # Which metric to optimize
+    direction: str = "maximize"  # "maximize" or "minimize"
+
+    # Output settings
+    output_dir: str = "sweep_results"
+
+    # Evaluation settings
+    skip_fid: bool = False
+
+    # Device configuration
+    device: str = "auto:1"
+
+    # Pruning settings
+    enable_pruning: bool = True
+    pruner_type: str = "median"  # "median", "hyperband", or "none"
+
+
+class OptunaObjective:
+    """Optuna objective function for hyperparameter optimization."""
+
+    def __init__(self, config: SweeperConfig):
+        self.config = config
+        self.base_config = self._load_base_config()
+        self.hazy_images, self.image_paths = load_images_from_dir(
+            config.input_image_dir
+        )
+
+        # Initialize device
+        init_device(config.device)
+
+        # Initialize the diffusion model once
+        self.diffusion_model = init(self.base_config)
+
+    def _load_base_config(self):
+        """Load base configuration from YAML file."""
+        with open(self.config.base_config_path, "r") as f:
+            config_dict = yaml.safe_load(f)
+        return zea.Config(**config_dict)
+
+    def _create_trial_params(self, trial: optuna.Trial) -> Dict[str, Any]:
+        """Create trial parameters by suggesting hyperparameters."""
+        params = {
+            "guidance_kwargs": {
+                "omega": trial.suggest_float("omega", 0.5, 50.0, log=True),
+                "omega_vent": trial.suggest_float("omega_vent", 0.0001, 50.0, log=True),
+                "omega_sept": trial.suggest_float("omega_sept", 0.1, 50.0, log=True),
+                "omega_dark": trial.suggest_float("omega_dark", 0.001, 50.0, log=True),
+                "eta": trial.suggest_float("eta", 0.001, 1.0, log=True),
+                "smooth_l1_beta": trial.suggest_float(
+                    "smooth_l1_beta", 0.1, 10.0, log=True
+                ),
+            },
+            "skeleton_params": {
+                "sigma_pre": trial.suggest_float("skeleton_sigma_pre", 0.0, 10.0),
+                "sigma_post": trial.suggest_float("skeleton_sigma_post", 0.0, 10.0),
+                "threshold": trial.suggest_float("skeleton_threshold", 0.0, 1.0),
+            },
+            "mask_params": {
+                "threshold": trial.suggest_float("mask_threshold", 0.0, 1.0),
+                "sigma": trial.suggest_float("mask_sigma", 0.0, 10.0),
+            },
+        }
+
+        # Add base config parameters that aren't being optimized
+        if hasattr(self.base_config, "params"):
+            base_params = self.base_config.params
+            for key, value in base_params.items():
+                if key not in params:
+                    params[key] = value
+
+        return params
+
+    def __call__(self, trial: optuna.Trial) -> float:
+        """Optuna objective function."""
+        # Suggest hyperparameters for this trial
+        params = self._create_trial_params(trial)
+
+        # Create seed for reproducibility
+        seed = jax.random.PRNGKey(self.base_config.seed + trial.number)
+
+        # Run the semantic DPS method
+        try:
+            hazy_images, pred_tissue_images, pred_haze_images, masks = run(
+                hazy_images=self.hazy_images,
+                diffusion_model=self.diffusion_model,
+                seed=seed,
+                **params,
+            )
+        except Exception as e:
+            log.error(f"Error during model inference: {e}")
+            return 0.0
+
+        # Convert tensors to numpy arrays if needed
+        if hasattr(pred_tissue_images, "numpy"):
+            pred_tissue_images = pred_tissue_images.numpy()
+
+        # Initialize temp directory
+        pred_tissue_temp_dir = None
+
+        try:
+            # Save predicted tissue images to temp directory
+            pred_tissue_temp_dir = save_images_to_temp_dir(
+                pred_tissue_images, self.image_paths, prefix="pred_tissue_"
+            )
+
+            # Run evaluation using the updated evaluate function
+            results = evaluate(
+                folder=pred_tissue_temp_dir,
+                noisy_folder=self.config.input_image_dir,
+                roi_folder=self.config.roi_folder,
+                reference_folder=self.config.reference_folder,
+            )
+
+            objective_value = results[self.config.objective_metric]
+
+        except Exception as e:
+            log.error(f"Error during evaluation: {e}")
+            objective_value = 0.0
+
+        finally:
+            # Clean up temporary directory
+            if pred_tissue_temp_dir and Path(pred_tissue_temp_dir).exists():
+                try:
+                    shutil.rmtree(pred_tissue_temp_dir)
+                except Exception as e:
+                    log.warning(
+                        f"Failed to clean up temp directory {pred_tissue_temp_dir}: {e}"
+                    )
+
+        # Log intermediate results for potential pruning
+        trial.report(objective_value, step=0)
+
+        # Check if trial should be pruned
+        if trial.should_prune():
+            raise optuna.TrialPruned()
+
+        # Store hyperparameters as user attributes
+        for key, value in params.items():
+            if isinstance(value, dict):
+                for subkey, subvalue in value.items():
+                    trial.set_user_attr(f"{key}_{subkey}", subvalue)
+            else:
+                trial.set_user_attr(key, value)
+
+        log.info(
+            f"Trial {trial.number}: {self.config.objective_metric} = {objective_value:.4f}"
+        )
+
+        return objective_value
+
+
+def create_pruner(pruner_type: str) -> optuna.pruners.BasePruner:
+    """Create an Optuna pruner based on the specified type."""
+    if pruner_type == "median":
+        return optuna.pruners.MedianPruner(
+            n_startup_trials=5, n_warmup_steps=0, interval_steps=1
+        )
+    elif pruner_type == "hyperband":
+        return optuna.pruners.HyperbandPruner(
+            min_resource=1, max_resource=100, reduction_factor=3
+        )
+    else:
+        return optuna.pruners.NopPruner()
+
+
+def run_optimization(config: SweeperConfig):
+    """Run hyperparameter optimization using Optuna."""
+
+    # Create pruner
+    pruner = create_pruner(config.pruner_type) if config.enable_pruning else None
+
+    # Create or load study
+    study = optuna.create_study(
+        study_name=config.study_name,
+        storage=config.storage,
+        direction=config.direction,
+        pruner=pruner,
+        load_if_exists=True,
+    )
+
+    log.info(f"Starting optimization for method: {config.method}")
+    log.info(f"Study name: {config.study_name}")
+    log.info(f"Number of trials: {config.n_trials}")
+    log.info(f"Objective metric: {config.objective_metric} ({config.direction})")
+
+    # Create objective function
+    objective = OptunaObjective(config)
+
+    # Run optimization
+    study.optimize(objective, n_trials=config.n_trials)
+
+    # Save results
+    output_dir = Path(config.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Save best trial info
+    best_trial = study.best_trial
+    best_results = {
+        "best_value": best_trial.value,
+        "best_params": best_trial.params,
+        "best_user_attrs": best_trial.user_attrs,
+        "study_stats": {
+            "n_trials": len(study.trials),
+            "n_complete_trials": len(
+                [t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]
+            ),
+            "n_pruned_trials": len(
+                [t for t in study.trials if t.state == optuna.trial.TrialState.PRUNED]
+            ),
+            "n_failed_trials": len(
+                [t for t in study.trials if t.state == optuna.trial.TrialState.FAIL]
+            ),
+        },
+    }
+
+    with open(
+        output_dir / f"best_results_{config.method}_{config.study_name}.json", "w"
+    ) as f:
+        json.dump(best_results, f, indent=2)
+
+    # Save all trials data
+    trials_data = []
+    for trial in study.trials:
+        trial_data = {
+            "number": trial.number,
+            "value": trial.value,
+            "params": trial.params,
+            "user_attrs": trial.user_attrs,
+            "state": trial.state.name,
+            "datetime_start": trial.datetime_start.isoformat()
+            if trial.datetime_start
+            else None,
+            "datetime_complete": trial.datetime_complete.isoformat()
+            if trial.datetime_complete
+            else None,
+        }
+        trials_data.append(trial_data)
+
+    with open(
+        output_dir / f"all_trials_{config.method}_{config.study_name}.json", "w"
+    ) as f:
+        json.dump(trials_data, f, indent=2)
+
+    # Print summary
+    log.success("Optimization completed!")
+    log.info(f"Best {config.objective_metric}: {best_trial.value:.4f}")
+    log.info("Best parameters:")
+    for key, value in best_trial.params.items():
+        log.info(f"  {key}: {value}")
+
+    # Print study statistics
+    stats = best_results["study_stats"]
+    log.info("Study statistics:")
+    log.info(f"  Total trials: {stats['n_trials']}")
+    log.info(f"  Complete trials: {stats['n_complete_trials']}")
+    log.info(f"  Pruned trials: {stats['n_pruned_trials']}")
+    log.info(f"  Failed trials: {stats['n_failed_trials']}")
+
+    return study
+
+
+def main():
+    """Main function for running hyperparameter optimization."""
+    config = tyro.cli(SweeperConfig)
+
+    # Validate required paths exist
+    required_paths = [
+        (config.input_image_dir, "Input image directory"),
+        (config.roi_folder, "ROI folder"),
+        (config.reference_folder, "Reference folder"),
+    ]
+
+    for path, description in required_paths:
+        if not Path(path).exists():
+            raise FileNotFoundError(f"{description} not found: {path}")
+
+    # Set visualization style
+    zea.visualize.set_mpl_style()
+
+    # Run optimization
+    study = run_optimization(config)
+
+    # Optionally, generate optimization plots
+    try:
+        import matplotlib.pyplot as plt
+        import optuna.visualization as vis
+
+        output_dir = Path(config.output_dir)
+
+        # Plot optimization history
+        fig = vis.matplotlib.plot_optimization_history(study).figure
+        fig.savefig(
+            output_dir / f"optimization_history_{config.method}.png",
+            dpi=300,
+            bbox_inches="tight",
+        )
+        plt.close(fig)
+
+        # Plot parameter importances
+        fig = vis.matplotlib.plot_param_importances(study).figure
+        fig.savefig(
+            output_dir / f"param_importances_{config.method}.png",
+            dpi=300,
+            bbox_inches="tight",
+        )
+        plt.close(fig)
+
+        # Plot parallel coordinate
+        fig = vis.matplotlib.plot_parallel_coordinate(study).figure
+        fig.savefig(
+            output_dir / f"parallel_coordinate_{config.method}.png",
+            dpi=300,
+            bbox_inches="tight",
+        )
+        plt.close(fig)
+
+        log.success(f"Optimization plots saved to {output_dir}")
+
+    except ImportError:
+        log.warning(
+            "Optuna visualization not available. Install with: pip install optuna[visualization]"
+        )
+
+
+if __name__ == "__main__":
+    main()