model/README_refinement.md

LeGR-Conv FiLM2 with BASNet Refinement Module

Overview

The legrconv4x_film2_refine.py model extends the FiLM2 architecture by integrating the refinement module (RefUnet) from BASNet (Boundary-Aware Salient Object Detection). This enhancement improves boundary accuracy for both segmentation masks and grasp predictions using an efficient two-module approach.

Key Components

1. RefUnet (from BASNet)

The refinement module directly adapted from BASNet:

• Takes coarse predictions as input (1 or 4 channels)
• Processes through encoder-decoder structure with skip connections
• Learns residual corrections
• Returns refined prediction by adding residual to input: output = input + residual

2. Multi-Task Adaptation

Our model uses an efficient two-module refinement approach:

• refine_mask: RefUnet for instance segmentation mask refinement (1-channel input/output)
• refine_grasp: RefUnet for joint grasp prediction refinement
  • Takes concatenated 4-channel input: [quality, sin, cos, width]
  • Learns correlations between grasp components
  • Outputs refined 4-channel predictions that are split back
  • Reduces model parameters from 5 to 2 RefUnet modules

3. Training Strategy

Following BASNet's multi-output supervision approach:

• Refined predictions: Primary loss with full weight (1.0)
• Coarse predictions: Auxiliary loss with weight 0.5
• This helps the network learn better representations at both stages

Architecture Details

RefUnet Structure (from BASNet)

Input (Coarse Prediction) [B, C, H, W]  # C=1 for mask, C=4 for grasp
    ↓
Conv0 → Conv1 → Pool1
    ↓      ↓(skip)
Conv2 → Pool2
    ↓      ↓(skip)
Conv3 → Pool3
    ↓      ↓(skip)
Conv4 → Pool4
    ↓      ↓(skip)
Conv5
    ↓
Upsample → Concat with Conv4 skip → Conv_d4
    ↓
Upsample → Concat with Conv3 skip → Conv_d3
    ↓
Upsample → Concat with Conv2 skip → Conv_d2
    ↓
Upsample → Concat with Conv1 skip → Conv_d1
    ↓
Conv_d0 (Residual) [B, C, H, W]
    ↓
Output = Input + Residual

Integration with FiLM2

1. Coarse Predictions: Generated through the standard FiLM2 pipeline
2. Refinement:
   • Mask refined independently through refine_mask
   • Grasp predictions concatenated and jointly refined through refine_grasp
3. Loss Computation: Both coarse and refined predictions are supervised

Usage

Training

python train_network_refactored.py \
    --config config/OCID-VLG/legrconv_film2_refine.yaml

Configuration

Key parameters in the config:

TRAIN:
  network_name: legrconv4x_film2_refine
  use_refinement: True  # Enable/disable refinement modules

Disabling Refinement (for ablation)

TRAIN:
  use_refinement: False  # Will use coarse predictions only

Implementation Highlights

Key Differences from Original BASNet

1. Two RefUnets: One for mask, one for joint grasp refinement (reduced from 5 modules)
2. Input Channels:
   • Mask RefUnet: 1-channel input/output
   • Grasp RefUnet: 4-channel input/output (concatenated predictions)
3. Auxiliary Loss Weight: 0.5 for coarse predictions (adjustable)

Loss Computation

# Refined prediction losses (primary)
loss_refined = compute_loss(refined_pred, target)

# Coarse prediction losses (auxiliary)
loss_coarse = 0.5 * compute_loss(coarse_pred, target)

# Total loss
total_loss = loss_refined + loss_coarse

Expected Improvements

Based on BASNet's performance:

• Better boundary accuracy: RefUnet focuses on boundary refinement
• Smoother predictions: Residual learning helps preserve details
• More stable training: Multi-level supervision improves convergence
• Efficient refinement: Joint grasp refinement leverages correlations between components
• Reduced parameters: Only 2 RefUnet modules instead of 5, making training more efficient

References

• BASNet: Qin et al., "BASNet: Boundary-Aware Salient Object Detection", CVPR 2019
• Original implementation: https://github.com/xuebinqin/BASNet