# 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
```bash
python train_network_refactored.py \
    --config config/OCID-VLG/legrconv_film2_refine.yaml
```

### Configuration
Key parameters in the config:
```yaml
TRAIN:
  network_name: legrconv4x_film2_refine
  use_refinement: True  # Enable/disable refinement modules
```

### Disabling Refinement (for ablation)
```yaml
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
```python
# 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