PanoFormer

https://github.com/zhijieshen-bjtu/panoformer
Admin
PanoFormer is a transformer-based deep learning model for indoor 360° panoramic depth estimation...
Tokens:4,313
Snippets:32
Trust Score:7.1
License:MIT
Update:6 months ago
Show doc for...
Tokens:
Raw
### Run Training Script

Source: https://github.com/zhijieshen-bjtu/panoformer/blob/main/README.md

Execute the main training script after downloading the pretrained model and placing it in the correct directory. Ensure all package dependencies are installed.

```python
python train.py
```

--------------------------------

### Install Package Dependencies

Source: https://github.com/zhijieshen-bjtu/panoformer/blob/main/README.md

Install all necessary Python packages listed in the requirements.txt file to set up the project environment. This is crucial for the project to run correctly.

```bash
pip install -r requirements.txt
```

--------------------------------

### Install Dependencies and Train PanoFormer

Source: https://context7.com/zhijieshen-bjtu/panoformer/llms.txt

Install necessary dependencies and train the PanoFormer model on Matterport3D or Stanford2D3D datasets using command-line arguments. Includes options for disabling data augmentation and monitoring training with TensorBoard.

```bash
pip install -r requirements.txt
```

```bash
python train.py \
    --batch_size 2 \
    --learning_rate 1e-4 \
    --num_epochs 200 \
    --gpu_devices 0 \
    --load_weights_dir ./tmp/panodepth/models/weights_pretrain
```

```bash
python trains2d3d.py \
    --batch_size 2 \
    --learning_rate 1e-4 \
    --num_epochs 200 \
    --gpu_devices 0 \
    --load_weights_dir ./tmp_s2d3d/panodepth/models/weights
```

```bash
python train.py \
    --disable_color_augmentation \
    --disable_LR_filp_augmentation \
    --disable_yaw_rotation_augmentation
```

```bash
tensorboard --logdir=./tmp/panodepth/
```

--------------------------------

### Run Stanford2D3D Training

Source: https://github.com/zhijieshen-bjtu/panoformer/blob/main/README.md

This command is used to initiate the training process specifically for the Stanford2D3D dataset. Ensure the dataset is downloaded and preprocessed correctly.

```python
python trains2d3d.py
```

--------------------------------

### Iterate through Training Batches

Source: https://context7.com/zhijieshen-bjtu/panoformer/llms.txt

Demonstrates how to access RGB, depth, and mask data from the training loader.

```python
for batch_idx, inputs in enumerate(train_loader):
    rgb = inputs["rgb"]                    # Original RGB (B, 3, 512, 1024)
    normalized_rgb = inputs["normalized_rgb"]  # Normalized for model input
    gt_depth = inputs["gt_depth"]          # Ground truth depth (B, 1, 512, 1024)
    val_mask = inputs["val_mask"]          # Valid depth mask

    print(f"RGB shape: {rgb.shape}")
    print(f"GT depth shape: {gt_depth.shape}")
    print(f"Valid mask shape: {val_mask.shape}")
    break
```

--------------------------------

### Initialize and Run PanoFormer Model

Source: https://context7.com/zhijieshen-bjtu/panoformer/llms.txt

Initializes the PanoFormer model with specified parameters and performs a forward pass with sample input. Ensure the model is moved to CUDA if a GPU is available. The output is a dictionary containing the predicted depth map.

```python
import torch
from network.model import Panoformer

# Initialize the model with default parameters
model = Panoformer(
    img_size=256,           # Base image size (actual input is 512x1024)
    in_chans=3,             # RGB input channels
    embed_dim=32,           # Base embedding dimension
    depths=[2, 2, 2, 2, 2, 2, 2, 2, 2],  # Depths for each encoder/decoder layer
    num_heads=[1, 2, 4, 8, 16, 16, 8, 4, 2],  # Attention heads per layer
    win_size=8,             # Window size for attention
    mlp_ratio=4.,           # MLP hidden dimension ratio
    drop_rate=0.,           # Dropout rate
    attn_drop_rate=0.,      # Attention dropout rate
    drop_path_rate=0.1      # Stochastic depth rate
)

# Move model to GPU
model = model.cuda()

# Forward pass with sample input (batch_size=2, channels=3, height=512, width=1024)
input_tensor = torch.randn(2, 3, 512, 1024).cuda()
outputs = model(input_tensor)

# Get predicted depth map (batch_size, 1, 512, 1024)
pred_depth = outputs["pred_depth"]
print(f"Predicted depth shape: {pred_depth.shape}")
# Output: Predicted depth shape: torch.Size([2, 1, 512, 1024])
```

--------------------------------

### Initialize Stanford2D3D Dataset

Source: https://context7.com/zhijieshen-bjtu/panoformer/llms.txt

Configures the dataset for panoramic RGB-D data with optional augmentations and pole masking.

```python
import torch
from torch.utils.data import DataLoader
from stanford2d3d import Stanford2D3D

# Initialize dataset
train_dataset = Stanford2D3D(
    root_dir='../data/Stanford2D3D/',
    list_file='./splits2d3d/stanford2d3d_train.txt',
    height=512,
    width=1024,
    disable_color_augmentation=False,
    disable_LR_filp_augmentation=False,
    disable_yaw_rotation_augmentation=False,
    is_training=True  # Training set is replicated 10x
)

val_dataset = Stanford2D3D(
    root_dir='../data/Stanford2D3D/',
    list_file='./splits2d3d/stanford2d3d_test.txt',
    is_training=False
)

# Create data loaders
train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=2, shuffle=False)

print(f"Training samples: {len(train_dataset)}")
print(f"Validation samples: {len(val_dataset)}")
```

--------------------------------

### Run Training Orchestrator

Source: https://context7.com/zhijieshen-bjtu/panoformer/llms.txt

Configures training arguments and initiates the training loop using the Trainer class.

```python
import argparse
from trainer import Trainer

# Setup training arguments
args = argparse.Namespace(
    num_workers=8,
    gpu_devices=[0],
    model_name="panodepth",
    learning_rate=1e-4,
    batch_size=2,
    num_epochs=200,
    load_weights_dir='./tmp/panodepth/models/weights_pretrain',
    log_dir='./tmp',
    log_frequency=100,
    save_frequency=1,
    disable_color_augmentation=False,
    disable_LR_filp_augmentation=False,
    disable_yaw_rotation_augmentation=False
)

# Initialize trainer
trainer = Trainer(args)

# Run training
trainer.train()  # Runs for num_epochs with validation after each epoch

# Model weights saved to: ./tmp/panodepth/models/weights_{epoch}/
```

--------------------------------

### Initialize and Use Panoramic Self-Attention (PSA)

Source: https://context7.com/zhijieshen-bjtu/panoformer/llms.txt

Initializes the PanoSelfAttention module, generates reference points, and performs a forward pass. This module is designed to handle equirectangular distortion in panoramic images using deformable attention.

```python
import torch
from network.PSA import PanoSelfAttention, generate_ref_points

# Initialize PSA module
psa = PanoSelfAttention(
    h=8,                    # Number of attention heads
    d_model=256,            # Model dimension
    k=9,                    # Number of sampling keys (3x3 pattern)
    last_feat_height=64,    # Feature map height
    last_feat_width=128,    # Feature map width
    scales=1,               # Number of scales
    dropout=0.1,
    need_attn=False         # Set True to return attention weights
).cuda()

# Generate reference points for sampling
ref_points = generate_ref_points(width=128, height=64)

# Input query and key tensors (B, H, W, C)
batch_size = 2
query = torch.randn(batch_size, 64, 128, 256).cuda()
keys = [query.unsqueeze(0)]  # List of key tensors for multi-scale

# Forward pass
ref_point_batch = ref_points.repeat(batch_size, 1, 1, 1)
output = psa(query, keys, ref_point_batch)
print(f"PSA output shape: {output.shape}")
# Output: PSA output shape: torch.Size([2, 64, 128, 256])
```

--------------------------------

### Load Matterport3D Dataset with Augmentations

Source: https://context7.com/zhijieshen-bjtu/panoformer/llms.txt

Initializes the Matterport3D dataset for training, enabling various data augmentations like color jittering, horizontal flipping, and yaw rotation. A DataLoader is then created for batch processing.

```python
import torch
from torch.utils.data import DataLoader
from matterport3d import Matterport3D

# Initialize dataset
train_dataset = Matterport3D(
    root_dir='../data/realM3D/',
    list_file='./splitsm3d/matterport3d_train.txt',
    height=512,
    width=1024,
    disable_color_augmentation=False,
    disable_LR_filp_augmentation=False,
    disable_yaw_rotation_augmentation=False,
    is_training=True
)

# Create data loader
train_loader = DataLoader(
    train_dataset,
    batch_size=2,
    shuffle=True,
    num_workers=8,
    pin_memory=True,
    drop_last=True
)

```

--------------------------------

### Evaluate Depth Metrics

Source: https://context7.com/zhijieshen-bjtu/panoformer/llms.txt

Computes standard depth estimation metrics and aggregates them across validation batches.

```python
import torch
from metrics import Evaluator, compute_depth_metrics

# Initialize evaluator
evaluator = Evaluator(median_align=True)

# Sample batch evaluation
gt_depth = torch.rand(4, 1, 512, 1024).cuda() * 10 + 0.5
pred_depth = gt_depth + torch.randn_like(gt_depth) * 0.5

# Compute metrics for single batch
mre, mae, abs_, abs_rel, sq_rel, rmse, rmse_log, log10, a1, a2, a3 = compute_depth_metrics(
    gt_depth, pred_depth
)

print(f"RMSE: {rmse.item():.4f}")
print(f"MAE: {mae.item():.4f}")
print(f"Delta1 (< 1.25): {a1.item():.4f}")
print(f"Delta2 (< 1.25^2): {a2.item():.4f}")
print(f"Delta3 (< 1.25^3): {a3.item():.4f}")

# Use evaluator for aggregated metrics across validation set
evaluator.reset_eval_metrics()
for i in range(10):  # Simulate multiple batches
    gt = torch.rand(2, 1, 512, 1024).cuda() * 10 + 0.5
    pred = gt + torch.randn_like(gt) * 0.3
    evaluator.compute_eval_metrics(gt, pred)

# Print aggregated results
evaluator.print()
```

--------------------------------

### PanoFormer Model Inference

Source: https://context7.com/zhijieshen-bjtu/panoformer/llms.txt

Performs inference to predict depth from a single panoramic image using a pretrained PanoFormer model. Includes image preprocessing, model loading, inference execution, and depth map post-processing for visualization.

```python
import torch
import cv2
import numpy as np
from torchvision import transforms
from network.model import Panoformer

# Load pretrained model
model = Panoformer()
model.cuda()
checkpoint = torch.load('./tmp/panodepth/models/weights_199/model.pth')
model.load_state_dict(checkpoint)
model.eval()

# Prepare input image
def preprocess_image(image_path):
    rgb = cv2.imread(image_path)
    rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
    rgb = cv2.resize(rgb, (1024, 512), interpolation=cv2.INTER_CUBIC)

    to_tensor = transforms.ToTensor()
    normalize = transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )

    tensor = normalize(to_tensor(rgb))
    return tensor.unsqueeze(0)  # Add batch dimension

# Run inference
input_tensor = preprocess_image('path/to/panorama.jpg').cuda()

with torch.no_grad():
    outputs = model(input_tensor)
    pred_depth = outputs["pred_depth"]

# Post-process depth map
depth_map = pred_depth.squeeze().cpu().numpy()
depth_map = np.clip(depth_map, 0.1, 10.0)  # Clip to valid range

# Save depth visualization
depth_normalized = (depth_map / depth_map.max() * 255).astype(np.uint8)
cv2.imwrite('predicted_depth.png', depth_normalized)
print(f"Depth range: {depth_map.min():.2f}m - {depth_map.max():.2f}m")
```

--------------------------------

### Generate Equirectangular Sampling Patterns

Source: https://context7.com/zhijieshen-bjtu/panoformer/llms.txt

Generates equirectangular sampling patterns for distortion-aware attention. These patterns account for non-uniform pixel distribution in panoramic images and are used in PanoformerBlock. The last dimension contains (latitude, longitude) offsets for each sampling point.

```python
import torch
from network.equisamplingpoint import genSamplingPattern

# Generate sampling pattern for different resolutions
# Parameters: height, width, kernel_height, kernel_width
pattern_256x512 = genSamplingPattern(256, 512, 3, 3).cuda()
pattern_128x256 = genSamplingPattern(128, 256, 3, 3).cuda()
pattern_64x128 = genSamplingPattern(64, 128, 3, 3).cuda()
pattern_32x64 = genSamplingPattern(32, 64, 3, 3).cuda()
pattern_16x32 = genSamplingPattern(16, 32, 3, 3).cuda()

print(f"256x512 pattern shape: {pattern_256x512.shape}")
print(f"128x256 pattern shape: {pattern_128x256.shape}")
```

--------------------------------

### Compute Berhu Loss

Source: https://context7.com/zhijieshen-bjtu/panoformer/llms.txt

Calculates the Berhu loss for depth estimation, optionally incorporating gradient-based loss terms.

```python
import torch
from losses import BerhuLoss, L1Loss, L2Loss

# Initialize loss functions
berhu_loss = BerhuLoss(threshold=0.2)  # Default threshold
l1_loss = L1Loss()
l2_loss = L2Loss()

# Sample predictions and ground truth
pred = torch.rand(2, 1, 512, 1024).cuda() * 10  # Predicted depth [0, 10] meters
target = torch.rand(2, 1, 512, 1024).cuda() * 10  # Ground truth depth

# Compute Berhu loss
loss = berhu_loss(target, pred)
print(f"Berhu loss: {loss.item():.4f}")

# For training with gradient loss
import loss_gradient as loss_g

gradient_net = loss_g.Gradient_Net()
gt_gx, gt_gy = gradient_net(target.float())
pred_gx, pred_gy = gradient_net(pred)

# Combined loss with gradient terms
total_loss = berhu_loss(target, pred) + berhu_loss(gt_gx, pred_gx) + berhu_loss(gt_gy, pred_gy)
print(f"Total loss with gradients: {total_loss.item():.4f}")
```
Tokens:
Raw