Initial commit

app.py

@@ -0,0 +1,228 @@
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib import cm
+import gradio as gr
+from PIL import Image
+from pathlib import Path
+
+# Import the inference module
+from inference import BathymetrySuperResolution
+
+# Define checkpoint and config paths
+CHECKPOINT_DIR = os.environ.get("CHECKPOINT_DIR", "checkpoints")
+MODEL_CHECKPOINT = os.path.join(CHECKPOINT_DIR, "calibrated.pth")
+CONFIG_PATH = os.environ.get("CONFIG_PATH", "config.json")
+
+# Initialize model
+try:
+    model = BathymetrySuperResolution(
+        model_type="vqvae",
+        checkpoint_path=MODEL_CHECKPOINT,
+        config_path=CONFIG_PATH
+    )
+    model_loaded = True
+except Exception as e:
+    print(f"Error loading model: {str(e)}")
+    model = None
+    model_loaded = False
+
+def process_upload(file, confidence_level, block_size, model_type):
+    """Process uploaded bathymetry file"""
+    if file is None:
+        return None, "Please upload a file."
+    
+    try:
+        # Check if the model is loaded
+        if not model_loaded:
+            return None, "Model not loaded. Please check server logs."
+        
+        # Load the data
+        if file.name.endswith('.npy'):
+            data = np.load(file.name)
+        else:
+            # Try to load as an image
+            img = Image.open(file.name).convert('L')
+            data = np.array(img)
+        
+        # Update model configuration if needed
+        if model.config['model_type'] != model_type or model.config['model_config']['block_size'] != block_size:
+            # In a real app, you would reload the model or adjust the configuration
+            pass
+        
+        # Run the prediction
+        prediction, lower_bound, upper_bound = model.predict(
+            data, 
+            with_uncertainty=True,
+            confidence_level=confidence_level/100.0  # Convert percentage to fraction
+        )
+        
+        # Calculate uncertainty width
+        uncertainty_width = model.get_uncertainty_width(lower_bound, upper_bound)
+        
+        # Create visualization
+        fig = plt.figure(figsize=(15, 10))
+        
+        # Original input (resized to 32x32 if needed)
+        ax1 = fig.add_subplot(231)
+        if data.shape != (32, 32):
+            from scipy.ndimage import zoom
+            zoom_factor = 32 / max(data.shape)
+            input_data = zoom(data, zoom_factor)
+        else:
+            input_data = data
+        im1 = ax1.imshow(input_data, cmap=cm.viridis)
+        ax1.set_title("Input (32x32)")
+        plt.colorbar(im1, ax=ax1)
+        
+        # Super-resolution output
+        ax2 = fig.add_subplot(232)
+        im2 = ax2.imshow(prediction[0, 0], cmap=cm.viridis)
+        ax2.set_title("Super-Resolution (64x64)")
+        plt.colorbar(im2, ax=ax2)
+        
+        # Lower bound
+        ax3 = fig.add_subplot(233)
+        im3 = ax3.imshow(lower_bound[0, 0], cmap=cm.viridis)
+        ax3.set_title(f"Lower Bound ({confidence_level}% CI)")
+        plt.colorbar(im3, ax=ax3)
+        
+        # Upper bound
+        ax4 = fig.add_subplot(234)
+        im4 = ax4.imshow(upper_bound[0, 0], cmap=cm.viridis)
+        ax4.set_title(f"Upper Bound ({confidence_level}% CI)")
+        plt.colorbar(im4, ax=ax4)
+        
+        # Uncertainty width visualization
+        ax5 = fig.add_subplot(235)
+        uncertainty_map = upper_bound[0, 0] - lower_bound[0, 0]
+        im5 = ax5.imshow(uncertainty_map, cmap='hot')
+        ax5.set_title("Uncertainty Width")
+        plt.colorbar(im5, ax=ax5)
+        
+        # 3D surface plot
+        ax6 = fig.add_subplot(236, projection='3d')
+        x = np.arange(0, prediction.shape[2])
+        y = np.arange(0, prediction.shape[3])
+        X, Y = np.meshgrid(x, y)
+        surf = ax6.plot_surface(X, Y, prediction[0, 0], cmap=cm.viridis, 
+                              linewidth=0, antialiased=True)
+        ax6.set_title("3D Bathymetry")
+        
+        plt.tight_layout()
+        
+        # Return the figure and a summary text
+        summary = f"""
+        **Super-Resolution Results:**
+        - **Model Type**: {model_type.upper()}
+        - **Block Size**: {block_size}×{block_size}
+        - **Confidence Level**: {confidence_level}%
+        - **Average Uncertainty Width**: {uncertainty_width:.4f}
+        - **Input Shape**: {data.shape}
+        - **Output Shape**: {prediction.shape[2:]}
+        """
+        
+        return fig, summary
+        
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        return None, f"Error processing file: {str(e)}"
+
+def create_sample_data():
+    """Create a sample bathymetry data file for demonstration"""
+    # Create a synthetic bathymetry profile with features
+    x = np.linspace(0, 1, 32)
+    y = np.linspace(0, 1, 32)
+    xx, yy = np.meshgrid(x, y)
+    
+    # Create a surface with a ridge and a valley
+    z = -4000 + 500 * np.sin(10 * xx) * np.cos(8 * yy) + 300 * np.exp(-((xx-0.3)**2 + (yy-0.7)**2)/0.1)
+    
+    # Save to a temporary file
+    sample_dir = Path("samples")
+    sample_dir.mkdir(exist_ok=True)
+    sample_path = sample_dir / "sample_bathymetry.npy"
+    np.save(sample_path, z)
+    
+    return str(sample_path)
+
+# Create the Gradio interface
+with gr.Blocks(title="Bathymetry Super-Resolution") as demo:
+    gr.Markdown("""
+    # Bathymetry Super-Resolution with Uncertainty Quantification
+    
+    This application demonstrates super-resolution of ocean floor (bathymetry) data with uncertainty estimates.
+    Upload a bathymetry file (NPY or image) to see the enhanced resolution output with confidence intervals.
+    
+    The model uses a **Vector Quantized Variational Autoencoder (VQ-VAE)** with **block-based uncertainty quantification**.
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            input_file = gr.File(label="Upload Bathymetry File (.npy or image)")
+            
+            with gr.Row():
+                confidence_level = gr.Slider(
+                    minimum=80, maximum=99, value=95, step=1,
+                    label="Confidence Level (%)"
+                )
+                
+                block_size = gr.Dropdown(
+                    choices=[1, 2, 4, 8, 64], value=4,
+                    label="Block Size"
+                )
+                
+                model_type = gr.Dropdown(
+                    choices=["vqvae", "srcnn", "gan"], value="vqvae",
+                    label="Model Type"
+                )
+            
+            with gr.Row():
+                process_btn = gr.Button("Generate Super-Resolution")
+                sample_btn = gr.Button("Load Sample Data")
+        
+        with gr.Column():
+            output_plots = gr.Plot(label="Super-Resolution Results")
+            output_text = gr.Markdown(label="Summary")
+    
+    # Set up button actions
+    process_btn.click(
+        fn=process_upload, 
+        inputs=[input_file, confidence_level, block_size, model_type], 
+        outputs=[output_plots, output_text]
+    )
+    
+    # Sample data generation
+    sample_btn.click(
+        fn=lambda: gr.update(value=create_sample_data()),
+        inputs=None,
+        outputs=input_file
+    )
+    
+    gr.Markdown("""
+    ## About This Model
+    
+    This model enhances the resolution of bathymetric data from 32×32 to 64×64 while providing uncertainty estimates.
+    It was trained on bathymetry data from multiple ocean regions including the Eastern Pacific Basin, Western Pacific Region, and Indian Ocean Basin.
+    
+    The uncertainty estimates help identify areas where the model is less confident in its predictions, which is crucial for:
+    - Risk assessment in coastal hazard modeling
+    - Climate change impact analysis
+    - Tsunami propagation simulation
+        
+    ## Model Performance
+    
+    | Model | SSIM | PSNR | MSE | MAE | UWidth | CalErr |
+    |-------|------|------|-----|-----|--------|--------|
+    | UA-VQ-VAE | 0.9433 | 26.8779 | 0.0021 | 0.0317 | 0.1046 | 0.0664 |
+    """)
+
+# Launch the demo
+if __name__ == "__main__":
+    if model_loaded:
+        print("Model loaded successfully. Starting Gradio interface.")
+    else:
+        print("Warning: Model not loaded. Demo will display errors when processing files.")
+        
+    demo.launch()

checkpoints/calibrated.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5c2ba632b12fe06e6c684c964ea07424b074e36ed533c3b03fa3bb4e8bf1c67e
+size 235336891

config.json

@@ -0,0 +1,16 @@
+{
+    "model_type": "vqvae",
+    "model_config": {
+      "in_channels": 1,
+      "hidden_dims": [32, 64, 128, 256],
+      "num_embeddings": 512,
+      "embedding_dim": 256,
+      "block_size": 4
+    },
+    "normalization": {
+      "mean": -3911.3894,
+      "std": 1172.8374,
+      "min": 0.0,
+      "max": 1.0
+    }
+  }

inference.py

@@ -0,0 +1,207 @@
+import os
+import torch
+import numpy as np
+from PIL import Image
+import torch.nn.functional as F
+import json
+
+# Import your model components
+from models.loader import ModelLoader
+from models.uncertainty import BlockUncertaintyTracker
+
+class BathymetrySuperResolution:
+    """
+    Bathymetry super-resolution model with uncertainty estimation
+    """
+    def __init__(self, model_type="vqvae", checkpoint_path=None, config_path=None):
+        """
+        Initialize the super-resolution model with uncertainty awareness
+        
+        Args:
+            model_type: Type of model ('srcnn', 'gan', or 'vqvae')
+            checkpoint_path: Path to model checkpoint
+            config_path: Path to configuration file
+        """
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+        # Load config if provided
+        if config_path is not None and os.path.exists(config_path):
+            with open(config_path, 'r') as f:
+                self.config = json.load(f)
+        else:
+            # Default configuration
+            self.config = {
+                "model_type": model_type,
+                "model_config": {
+                    "in_channels": 1,
+                    "hidden_dims": [32, 64, 128, 256],
+                    "num_embeddings": 512,
+                    "embedding_dim": 256,
+                    "block_size": 4
+                },
+                "normalization": {
+                    "mean": -3911.3894,
+                    "std": 1172.8374,
+                    "min": 0.0,
+                    "max": 1.0
+                }
+            }
+        
+        # Initialize model loader
+        self.model_loader = ModelLoader()
+        
+        # Load model
+        if checkpoint_path is not None and os.path.exists(checkpoint_path):
+            self.model = self.model_loader.load_model(
+                self.config['model_type'],
+                checkpoint_path,
+                config_overrides=self.config.get('model_config', {})
+            )
+        else:
+            raise ValueError("Checkpoint path not provided or invalid")
+        
+        # Ensure model is in eval mode
+        self.model.eval()
+        
+        # Load normalization parameters
+        self.mean = self.config['normalization']['mean']
+        self.std = self.config['normalization']['std']
+        self.min_val = self.config['normalization']['min']
+        self.max_val = self.config['normalization']['max']
+    
+    def preprocess(self, data):
+        """
+        Preprocess input data for the model
+        
+        Args:
+            data: Input array/image (can be numpy array, PIL Image, or tensor)
+            
+        Returns:
+            Preprocessed tensor
+        """
+        # Convert PIL Image to numpy if needed
+        if isinstance(data, Image.Image):
+            data = np.array(data)
+        
+        # Convert numpy to tensor if needed
+        if isinstance(data, np.ndarray):
+            tensor = torch.from_numpy(data).float()
+        else:
+            tensor = data.float()
+        
+        # Add batch and channel dimensions if needed
+        if len(tensor.shape) == 2:
+            tensor = tensor.unsqueeze(0).unsqueeze(0)
+        elif len(tensor.shape) == 3:
+            tensor = tensor.unsqueeze(0)
+        
+        # Apply normalization
+        tensor = (tensor - self.mean) / (self.std + 1e-8)
+        tensor = (tensor - tensor.min()) / (tensor.max() - tensor.min() + 1e-8)
+        
+        # Resize if needed (to 32x32)
+        if tensor.shape[-1] != 32 or tensor.shape[-2] != 32:
+            tensor = F.interpolate(
+                tensor,
+                size=(32, 32),
+                mode='bicubic',
+                align_corners=False
+            )
+        
+        return tensor.to(self.device)
+    
+    def denormalize(self, tensor):
+        """
+        Denormalize output tensor
+        
+        Args:
+            tensor: Output tensor from model
+            
+        Returns:
+            Denormalized tensor in original data range
+        """
+        # Scale from [0,1] back to original range
+        tensor = tensor * (self.max_val - self.min_val) + self.min_val
+        
+        # Restore original scale
+        tensor = tensor * self.std + self.mean
+        
+        return tensor
+    
+    def predict(self, data, with_uncertainty=True, confidence_level=0.95):
+        """
+        Generate super-resolution output with uncertainty bounds
+        
+        Args:
+            data: Input data (can be numpy array, PIL Image, or tensor)
+            with_uncertainty: Whether to include uncertainty bounds
+            confidence_level: Confidence level for uncertainty bounds
+            
+        Returns:
+            Tuple of (prediction, lower_bound, upper_bound) if with_uncertainty=True
+            or just prediction otherwise
+        """
+        # Preprocess input
+        input_tensor = self.preprocess(data)
+        
+        with torch.no_grad():
+            # Run model inference
+            if with_uncertainty and hasattr(self.model, 'predict_with_uncertainty'):
+                prediction, lower_bound, upper_bound = self.model.predict_with_uncertainty(
+                    input_tensor, confidence_level
+                )
+                
+                # Denormalize outputs
+                prediction = self.denormalize(prediction)
+                lower_bound = self.denormalize(lower_bound) if lower_bound is not None else None
+                upper_bound = self.denormalize(upper_bound) if upper_bound is not None else None
+                
+                # Convert to numpy
+                prediction = prediction.cpu().numpy()
+                lower_bound = lower_bound.cpu().numpy() if lower_bound is not None else None
+                upper_bound = upper_bound.cpu().numpy() if upper_bound is not None else None
+                
+                return prediction, lower_bound, upper_bound
+            else:
+                # Standard inference
+                prediction = self.model(input_tensor)
+                
+                # Denormalize
+                prediction = self.denormalize(prediction)
+                
+                # Convert to numpy
+                prediction = prediction.cpu().numpy()
+                
+                return prediction
+    
+    def load_npy(self, file_path):
+        """
+        Load bathymetry data from numpy file
+        
+        Args:
+            file_path: Path to .npy file
+            
+        Returns:
+            Numpy array containing bathymetry data
+        """
+        try:
+            return np.load(file_path)
+        except Exception as e:
+            raise ValueError(f"Error loading numpy file: {str(e)}")
+    
+    @staticmethod
+    def get_uncertainty_width(lower_bound, upper_bound):
+        """
+        Calculate uncertainty width (difference between upper and lower bounds)
+        
+        Args:
+            lower_bound: Lower uncertainty bound
+            upper_bound: Upper uncertainty bound
+            
+        Returns:
+            Uncertainty width
+        """
+        if lower_bound is None or upper_bound is None:
+            return None
+        
+        return np.mean(upper_bound - lower_bound)

models/loader.py

@@ -0,0 +1,147 @@
+import torch
+import torch.nn as nn
+from enum import Enum
+from typing import Dict, Optional
+
+class ModelType(Enum):
+    SRCNN = 'srcnn'
+    GAN = 'gan'
+    VQVAE = 'vqvae'
+
+class ModelLoader:
+    """
+    Loader for different super-resolution model architectures
+    """
+    def __init__(self):
+        # Base model configurations
+        self.model_configs = {
+            ModelType.SRCNN: {    
+                "in_channels": 1,  
+                "hidden_channels": 64,
+                "num_residual_blocks": 8,
+                "num_upsamples": 1,          
+                "block_size": 4              
+            },
+            ModelType.GAN: {
+                "in_channels": 1,
+                "hidden_channels": 64,
+                "num_rrdb_blocks": 8,
+                "growth_channels": 32,
+                "num_upsamples": 1,          
+                "block_size": 4              
+            },
+            ModelType.VQVAE: {
+                "in_channels": 1,
+                "hidden_dims": [32, 64, 128, 256],
+                "num_embeddings": 512,
+                "embedding_dim": 256,
+                "block_size": 4
+            }
+        }
+        
+        self.model_registry = {}
+        self.loaded_models = {}
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+        # Try to import model classes
+        try:
+            from .vqvae import VQVAE
+            self.model_registry[ModelType.VQVAE] = VQVAE
+        except ImportError:
+            print("Warning: VQVAE model implementation not found")
+            
+        try:
+            from .cnn import CNN
+            self.model_registry[ModelType.SRCNN] = CNN
+        except ImportError:
+            print("Warning: CNN model implementation not found")
+            
+        try:
+            from .gan import UncertainESRGAN
+            self.model_registry[ModelType.GAN] = UncertainESRGAN
+        except ImportError:
+            print("Warning: GAN model implementation not found")
+
+    def load_model(self, model_type: str, checkpoint_path: str, config_overrides: Optional[Dict] = None):
+        """
+        Load a model with its checkpoint and optional configuration overrides
+        
+        Args:
+            model_type: Type of model to load ('srcnn', 'gan', or 'vqvae')
+            checkpoint_path: Path to model checkpoint
+            config_overrides: Optional dictionary of configuration overrides
+            
+        Returns:
+            Loaded model or None if loading fails
+        """
+        try:
+            # Convert string to enum
+            model_type = ModelType(model_type.lower())
+            
+            # Check if model implementation is available
+            if model_type not in self.model_registry:
+                raise ValueError(f"Model type {model_type.value} is not available")
+            
+            # Get base config and apply overrides if provided
+            model_config = self.model_configs[model_type].copy()
+            if config_overrides:
+                model_config.update(config_overrides)
+            
+            # Initialize model with potentially modified config
+            model_class = self.model_registry[model_type]
+            model = model_class(**model_config)
+            
+            # Move model to device
+            model = model.to(self.device)
+
+            # Load checkpoint
+            try:
+                checkpoint = torch.load(checkpoint_path, map_location=self.device)
+                model.load_state_dict(checkpoint['state_dict'], strict=False)
+                
+                # Load uncertainty tracker state if available
+                if hasattr(model, 'uncertainty_tracker'):
+                    model.uncertainty_tracker.calibrated = checkpoint.get('calibrated', False)
+                    
+                    if model.uncertainty_tracker.calibrated:
+                        if 'block_scale_means' in checkpoint:
+                            model.uncertainty_tracker.block_scale_means = checkpoint['block_scale_means']
+                        if 'block_scale_stds' in checkpoint:
+                            model.uncertainty_tracker.block_scale_stds = checkpoint['block_scale_stds']
+                
+                print(f"Successfully loaded {model_type.value} model from {checkpoint_path}")
+            except Exception as e:
+                print(f"Warning: Could not load checkpoint. Using untrained model. Error: {e}")
+
+            # Store model
+            model.eval()
+            self.loaded_models[model_type] = model
+            
+            return model
+
+        except Exception as e:
+            print(f"Error loading model: {str(e)}")
+            return None
+
+    def get_model(self, model_type: str):
+        """Get a loaded model by type"""
+        try:
+            model_type = ModelType(model_type.lower())
+            return self.loaded_models.get(model_type)
+        except:
+            return None
+
+    def unload_model(self, model_type: str):
+        """Unload a specific model"""
+        try:
+            model_type = ModelType(model_type.lower())
+            if model_type in self.loaded_models:
+                del self.loaded_models[model_type]
+                torch.cuda.empty_cache()
+        except:
+            pass
+
+    def unload_all_models(self):
+        """Unload all loaded models"""
+        self.loaded_models.clear()
+        torch.cuda.empty_cache()

models/uncertainty.py

@@ -0,0 +1,188 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class BlockUncertaintyTracker(nn.Module):
+    """
+    Track and estimate uncertainty at block level for bathymetry super-resolution
+    """
+    def __init__(self, block_size=4, alpha=0.1, decay=0.99, eps=1e-5):
+        """
+        Initialize block-wise uncertainty tracker
+        
+        Args:
+            block_size: Size of spatial blocks for uncertainty estimation
+            alpha: Quantile parameter for uncertainty bounds
+            decay: EMA decay factor for tracking statistics
+            eps: Small value for numerical stability
+        """
+        super().__init__()
+        self.block_size = block_size
+        self.decay = decay
+        self.alpha = alpha
+        self.eps = eps
+        
+        # Initialize unfold layer for block extraction
+        self.unfold = nn.Unfold(kernel_size=block_size, stride=block_size)
+        
+        # Register buffers with initial values
+        self.register_buffer('ema_errors', None)
+        self.register_buffer('ema_quantile', None)
+        
+        self.num_blocks_h = None
+        self.num_blocks_w = None
+
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+        # Calibration statistics
+        self.calibrated = False
+        self.block_means = []
+        self.block_stds = []
+        self.block_scale_means = None
+        self.block_scale_stds = None
+
+    def _initialize_buffers(self, h, w, device):
+        """Initialize EMA buffers based on number of blocks in image"""
+        self.num_blocks_h = h // self.block_size
+        self.num_blocks_w = w // self.block_size
+        num_blocks = self.num_blocks_h * self.num_blocks_w
+        
+        # Initialize buffers on the correct device
+        self.ema_errors = torch.zeros(num_blocks, device=device)
+        self.ema_quantile = torch.zeros(num_blocks, device=device)
+
+    def update(self, current_errors):
+        """Update EMA of errors and quantiles for each block"""
+        B, C, H, W = current_errors.shape
+        device = current_errors.device
+        
+        # Initialize buffers if not done yet
+        if self.ema_errors is None:
+            self._initialize_buffers(H, W, device)
+            
+        # Unfold into blocks
+        blocks = self.unfold(current_errors)
+        block_errors = blocks.transpose(1, 2)
+        block_errors = block_errors.reshape(-1, self.num_blocks_h * self.num_blocks_w, self.block_size * self.block_size)
+        
+        with torch.no_grad():
+            # Compute mean error per block
+            block_mean_errors = block_errors.mean(dim=-1)  # [B, num_blocks]
+            
+            # Update EMA errors for each block
+            block_means = block_mean_errors.mean(dim=0)  # Average across batch
+            block_means = block_means.to(device)  # Ensure on correct device
+            self.ema_errors = self.ema_errors.to(device)  # Ensure on correct device
+            self.ema_errors.mul_(self.decay).add_(block_means * (1 - self.decay))
+            
+            # Update quantiles for each block
+            block_quantiles = torch.quantile(block_errors, 1 - self.alpha, dim=-1)  # [B, num_blocks]
+            quantile_means = block_quantiles.mean(dim=0)  # Average across batch
+            quantile_means = quantile_means.to(device)  # Ensure on correct device
+            self.ema_quantile = self.ema_quantile.to(device)  # Ensure on correct device
+            self.ema_quantile.mul_(self.decay).add_(quantile_means * (1 - self.decay))
+
+    def get_uncertainty(self, errors):
+        """Calculate block-wise uncertainty scores"""
+        B, C, H, W = errors.shape
+        device = errors.device
+        
+        # Initialize buffers if not done yet
+        if self.ema_errors is None:
+            self._initialize_buffers(H, W, device)
+            
+        # Ensure buffers are on correct device
+        self.ema_errors = self.ema_errors.to(device)
+        self.ema_quantile = self.ema_quantile.to(device)
+        
+        # Unfold into blocks
+        blocks = self.unfold(errors)
+        
+        # Calculate uncertainty for each block
+        uncertainties = []
+        for i in range(self.num_blocks_h * self.num_blocks_w):
+            block = blocks[:, :, i].view(B, C, self.block_size, self.block_size)
+            uncertainty = block / (self.ema_quantile[i] + self.eps)
+            uncertainties.append(uncertainty)
+            
+        # Reconstruct full image from blocks
+        uncertainty_blocks = torch.stack(uncertainties, dim=-1)
+        uncertainty_blocks = uncertainty_blocks.permute(0, 1, 4, 2, 3)
+        
+        # Reshape to original image size
+        uncertainty_map = uncertainty_blocks.reshape(
+            B, C,
+            self.num_blocks_h, self.block_size,
+            self.num_blocks_w, self.block_size
+        ).permute(0, 1, 2, 4, 3, 5).reshape(B, C, H, W)
+        
+        return uncertainty_map
+
+    def get_bounds(self, x, confidence_level=0.95):
+        """
+        Get prediction bounds based on calibrated statistics
+        
+        Args:
+            x: Input tensor [B, C, H, W]
+            confidence_level: Confidence level for bounds
+            
+        Returns:
+            tuple: (lower_bounds, upper_bounds)
+        """
+        if not self.calibrated:
+            print("Warning: Model not calibrated. Bounds may be inaccurate.")
+            # Return simple bounds based on mean error
+            return x * 0.9, x * 1.1
+
+        # Calculate z-score based on confidence level
+        z_scores = {
+            0.99: 2.576,
+            0.95: 1.96,
+            0.90: 1.645,
+            0.85: 1.440,
+            0.80: 1.282
+        }
+        z_score = z_scores.get(confidence_level, 1.96)
+
+        # Get block-wise uncertainty
+        blocks = self.unfold(x)  # [B, C*block_size*block_size, num_blocks]
+        
+        # Calculate bounds for each block using calibrated statistics
+        B, C, H, W = x.shape
+        lower_bounds = []
+        upper_bounds = []
+        
+        for i in range(self.num_blocks_h * self.num_blocks_w):
+            block = blocks[:, :, i].view(B, C, self.block_size, self.block_size)
+            
+            # Use calibrated statistics to determine uncertainty
+            if hasattr(self, 'block_scale_stds') and self.block_scale_stds is not None:
+                uncertainty = z_score * self.block_scale_stds[i]
+            else:
+                # Fallback if calibration stats not available
+                uncertainty = 0.1 * block.mean()
+            
+            lower_bound = torch.clamp(block - uncertainty, min=0.0)
+            upper_bound = torch.clamp(block + uncertainty, max=1.0)
+            
+            lower_bounds.append(lower_bound)
+            upper_bounds.append(upper_bound)
+        
+        # Reconstruct full image from blocks
+        lower_bounds = torch.stack(lower_bounds, dim=-1)
+        upper_bounds = torch.stack(upper_bounds, dim=-1)
+        
+        # Reshape to original image size
+        lower_bounds = lower_bounds.permute(0, 1, 4, 2, 3).reshape(
+            B, C,
+            self.num_blocks_h, self.num_blocks_w,
+            self.block_size, self.block_size
+        ).permute(0, 1, 2, 4, 3, 5).reshape(B, C, H, W)
+        
+        upper_bounds = upper_bounds.permute(0, 1, 4, 2, 3).reshape(
+            B, C,
+            self.num_blocks_h, self.num_blocks_w,
+            self.block_size, self.block_size
+        ).permute(0, 1, 2, 4, 3, 5).reshape(B, C, H, W)
+        
+        return lower_bounds, upper_bounds

models/vqvae.py

@@ -0,0 +1,262 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .uncertainty import BlockUncertaintyTracker
+
+class ResidualAttentionBlock(nn.Module):
+    """Residual attention block for capturing spatial dependencies"""
+    def __init__(self, in_channels):
+        super().__init__()
+
+        # Trunk branch
+        self.trunk = nn.Sequential(
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=0),
+            nn.BatchNorm2d(in_channels),
+            nn.SiLU(),
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=0),
+            nn.BatchNorm2d(in_channels)
+        )
+
+        # Mask branch for attention
+        self.mask = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(in_channels, in_channels, kernel_size=1),
+            nn.SiLU(),
+            nn.Conv2d(in_channels, in_channels, kernel_size=1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        # Trunk branch
+        trunk_output = self.trunk(x)
+        
+        # Mask branch for attention weights
+        attention = self.mask(x)
+        
+        # Apply attention and residual connection
+        out = x + attention * trunk_output
+        return F.silu(out)
+
+class VectorQuantizer(nn.Module):
+    """Vector quantizer for discrete latent representation"""
+    def __init__(self, n_embeddings=512, embedding_dim=256, beta=0.25):
+        super().__init__()
+        self.n_embeddings = n_embeddings
+        self.embedding_dim = embedding_dim
+        self.beta = beta
+        
+        # Initialize embeddings
+        self.embeddings = nn.Parameter(torch.randn(n_embeddings, embedding_dim))
+        nn.init.uniform_(self.embeddings, -1.0 / n_embeddings, 1.0 / n_embeddings)
+        
+        # Usage tracking
+        self.register_buffer('usage', torch.zeros(n_embeddings))
+
+    def forward(self, z):
+        # Reshape input for quantization
+        z_flattened = z.reshape(-1, self.embedding_dim)
+        
+        # Calculate distances to embedding vectors
+        distances = torch.sum(z_flattened**2, dim=1, keepdim=True) + \
+                    torch.sum(self.embeddings**2, dim=1) - \
+                    2 * torch.matmul(z_flattened, self.embeddings.t())
+                    
+        # Find nearest embedding for each input vector
+        encoding_indices = torch.argmin(distances, dim=1)
+        
+        # Update usage statistics
+        if self.training:
+            with torch.no_grad():
+                usage = torch.zeros_like(self.usage)
+                usage.scatter_add_(0, encoding_indices, torch.ones_like(encoding_indices, dtype=torch.float))
+                self.usage.mul_(0.99).add_(usage, alpha=0.01)
+        
+        # Get quantized vectors
+        z_q = self.embeddings[encoding_indices].reshape(z.shape)
+        
+        # Calculate loss terms
+        commitment_loss = F.mse_loss(z_q.detach(), z)
+        codebook_loss = F.mse_loss(z_q, z.detach())
+        
+        # Combine losses
+        loss = codebook_loss + self.beta * commitment_loss
+        
+        # Straight-through estimator
+        z_q = z + (z_q - z).detach()
+        
+        if self.training:
+            return z_q, loss
+        else:
+            return z_q
+
+class Encoder(nn.Module):
+    """Encoder for VQ-VAE model"""
+    def __init__(self, in_channels=1, hidden_dims=[32, 64, 128, 256], embedding_dim=256):
+        super().__init__()
+        
+        # Initial conv layer
+        layers = [
+            nn.Conv2d(in_channels, hidden_dims[0], kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(hidden_dims[0]),
+            nn.SiLU()
+        ]
+        
+        # Hidden layers with downsampling
+        for i in range(len(hidden_dims) - 1):
+            layers.extend([
+                nn.Conv2d(hidden_dims[i], hidden_dims[i+1], kernel_size=4, stride=2, padding=1),
+                nn.BatchNorm2d(hidden_dims[i+1]),
+                nn.SiLU()
+            ])
+        
+        # Residual attention blocks
+        for _ in range(2):
+            layers.append(ResidualAttentionBlock(hidden_dims[-1]))
+        
+        # Final projection to embedding dimension
+        layers.extend([
+            nn.Conv2d(hidden_dims[-1], embedding_dim, kernel_size=1),
+            nn.BatchNorm2d(embedding_dim)
+        ])
+        
+        self.encoder = nn.Sequential(*layers)
+        
+    def forward(self, x):
+        return self.encoder(x)
+
+class Decoder(nn.Module):
+    """Decoder for VQ-VAE model"""
+    def __init__(self, embedding_dim=256, hidden_dims=[256, 128, 64, 32], out_channels=1):
+        super().__init__()
+        
+        # Reverse hidden dims for decoder
+        hidden_dims = hidden_dims[::-1]
+        
+        # Initial processing
+        layers = [
+            nn.Conv2d(embedding_dim, hidden_dims[0], kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(hidden_dims[0]),
+            nn.SiLU()
+        ]
+        
+        # Residual attention blocks
+        for _ in range(2):
+            layers.append(ResidualAttentionBlock(hidden_dims[0]))
+        
+        # Upsampling blocks
+        for i in range(len(hidden_dims) - 1):
+            layers.extend([
+                nn.ConvTranspose2d(hidden_dims[i], hidden_dims[i+1], 
+                                  kernel_size=4, stride=2, padding=1),
+                nn.BatchNorm2d(hidden_dims[i+1]),
+                nn.SiLU()
+            ])
+        
+        # Final output layer
+        layers.append(
+            nn.Conv2d(hidden_dims[-1], out_channels, kernel_size=3, padding=1)
+        )
+        layers.append(nn.Sigmoid())
+        
+        self.decoder = nn.Sequential(*layers)
+        
+    def forward(self, x):
+        return self.decoder(x)
+
+class VQVAE(nn.Module):
+    """
+    Vector Quantized Variational Autoencoder with uncertainty awareness
+    for bathymetry super-resolution
+    """
+    def __init__(self, in_channels=1, hidden_dims=[32, 64, 128, 256],
+                 num_embeddings=512, embedding_dim=256, block_size=4, alpha=0.1):
+        super().__init__()
+        
+        # Initialize block-wise uncertainty tracking
+        self.uncertainty_tracker = BlockUncertaintyTracker(
+            block_size=block_size,
+            alpha=alpha,
+            decay=0.99,
+            eps=1e-5
+        )
+        
+        # Main model components
+        self.encoder = Encoder(
+            in_channels=in_channels,
+            hidden_dims=hidden_dims,
+            embedding_dim=embedding_dim
+        )
+        
+        self.vq = VectorQuantizer(
+            n_embeddings=num_embeddings,
+            embedding_dim=embedding_dim,
+            beta=0.25
+        )
+        
+        self.decoder = Decoder(
+            embedding_dim=embedding_dim,
+            hidden_dims=hidden_dims,
+            out_channels=in_channels
+        )
+        
+    def forward(self, x):
+        """Forward pass through the model"""
+        # Encode
+        z = self.encoder(x)
+        
+        # Vector quantization
+        if self.training:
+            z_q, vq_loss = self.vq(z)
+            
+            # Decode
+            reconstruction = self.decoder(z_q)
+            
+            return reconstruction, vq_loss
+        else:
+            z_q = self.vq(z)
+            
+            # Decode
+            reconstruction = self.decoder(z_q)
+            
+            return reconstruction
+    
+    def train_forward(self, x, y):
+        """Training forward pass with uncertainty tracking"""
+        # Get reconstruction and VQ loss
+        reconstruction, vq_loss = self.forward(x)
+        
+        # Calculate reconstruction error
+        error = torch.abs(reconstruction - y)
+        
+        # Update uncertainty tracker
+        self.uncertainty_tracker.update(error)
+        
+        # Get uncertainty map for loss weighting
+        uncertainty_map = self.uncertainty_tracker.get_uncertainty(error)
+        
+        return reconstruction, vq_loss, uncertainty_map
+    
+    def predict_with_uncertainty(self, x, confidence_level=0.95):
+        """
+        Forward pass with calibrated uncertainty bounds
+        
+        Args:
+            x: Input tensor
+            confidence_level: Confidence level for bounds (default: 0.95)
+            
+        Returns:
+            tuple: (reconstruction, lower_bounds, upper_bounds)
+        """
+        self.eval()
+        with torch.no_grad():
+            # Get reconstruction
+            reconstruction = self.forward(x)
+            
+            # Get calibrated uncertainty bounds
+            lower_bounds, upper_bounds = self.uncertainty_tracker.get_bounds(
+                reconstruction, confidence_level
+            )
+            
+            return reconstruction, lower_bounds, upper_bounds

requirements.txt

@@ -0,0 +1,11 @@
+torch>=2.0.0
+numpy>=1.20.0
+matplotlib>=3.5.0
+gradio>=3.32.0
+Pillow>=9.0.0
+scipy>=1.8.0
+tqdm>=4.62.0
+huggingface_hub>=0.14.0
+transformers>=4.30.0
+pandas>=1.3.0
+scikit-learn>=1.0.0