You can deploy this model and start earning money today!

longvo2003t/
FM_PhysMamba_UNET

Name: FM_PhysMamba_UNET
Availability: InStock

Views

No views yet

FM_PhysMamba_UNET: Physics-Informed Mamba for Image Dehazing

This repository contains the official implementation of FM_PhysMamba_UNET, a State Space Model (Mamba) architecture integrated with Flow Matching and physical priors for high-performance image dehazing.

Image Illustrations

Figure 1: Comparison between Hazy Input (Left) and FM-PhysMamba Restored Output (Right)

📂 Repository Structure

To ensure all internal imports (e.g., from model import ...) function correctly, please maintain this layout:

model/: Contains the core UNet and PhysMamba blocks.
configs/: Model configuration files.
utils.py: Contains the utility logic for model inference.
data/: Data processing and re-standardization utilities.
checkpoints/: Directory for model weights (pytorch_model.bin) (Please put your weights inside this directory)
test_imgs/: Sample images for testing.

🚀 Quick Start: Inference

This script performs dehazing on large images by breaking them into overlapping tiles to manage memory efficiently.

Python

import torch
import os
from PIL import Image
from model import FM_PhysMamba_UNET, ODESolver
from utils import predict_large_image_vectorized, preprocess_single_image
from data.utils import restandardize_tensor
import matplotlib.pyplot as plt

# ==========================================
# CONFIGURATION
# ==========================================
# REPLACE THIS with your actual weight file path
WEIGHT_PATH = "checkpoints/final_weights/DENSE_HAZE/pytorch_model.bin" 

INPUT_IMG_PATH = "test_imgs/test_images.jpeg"
OUTPUT_DIR = "test_imgs/results"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

def load_dehazing_model(weight_path, device):
    """Initializes model and loads pre-trained weights."""
    print(f"Loading model to {device}...")
    model = FM_PhysMamba_UNET("small").to(device)
    
    if not os.path.exists(weight_path):
        raise FileNotFoundError(f"Weights not found at: {weight_path}")
        
    checkpoint = torch.load(weight_path, map_location=device, weights_only=True)
    model.load_state_dict(checkpoint)
    model.eval()
    return model

def run_inference(model, img_path, device, output_dir):
    """Performs tiled inference and saves comparison result."""
    solver = ODESolver(model)
    raw_img = Image.open(img_path).convert("RGB")
    input_tensor = preprocess_single_image(raw_img, device=device)
    
    print(f"Starting tiled inference for {img_path}...")
    with torch.no_grad():
        restored_tensor = predict_large_image_vectorized(
            solver=solver, 
            full_img_tensor=input_tensor, 
            device=device,
            tile_size=256,
            overlap_ratio=0.25
        )
    
    # Process for visualization
    hazy_disp = restandardize_tensor(input_tensor.detach().squeeze(0).cpu()).permute(1, 2, 0).numpy()
    restored_disp = restandardize_tensor(restored_tensor.detach().squeeze(0).cpu()).permute(1, 2, 0).numpy()

    # Create Comparison Plot
    os.makedirs(output_dir, exist_ok=True)
    fig, axes = plt.subplots(1, 2, figsize=(15, 7))
    axes[0].imshow(hazy_disp)
    axes[0].set_title("Original Hazy Input")
    axes[0].axis("off")
    
    axes[1].imshow(restored_disp)
    axes[1].set_title("FM-PhysMamba Restored")
    axes[1].axis("off")
    
    save_path = os.path.join(output_dir, "comparison_result.png")
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    print(f"Success! Result saved to: {save_path}")
    plt.show()

if __name__ == "__main__":
    # 1. Load the model once
    my_model = load_dehazing_model(WEIGHT_PATH, DEVICE)
    
    # 2. Run inference
    run_inference(my_model, INPUT_IMG_PATH, DEVICE, OUTPUT_DIR)

✨ Key Features

PhysMamba Block: Combines the efficiency of State Space Models (SSMs) with physics-based architectural constraints.
Flow Matching: Uses continuous-time normalizing flows for high-fidelity image restoration.
Vectorized Tiling: The predict_large_image_vectorized function in utils.py handles high-resolution inputs (2K/4K) without running out of VRAM.

🛠 Installation (using uv)

We recommend using uv for extremely fast dependency resolution.

Bash

# Install dependencies into a new virtual environment
uv sync

# Activate the environment
source .venv/bin/activate

longvo2003t/FM_PhysMamba_UNET