CBAM-STN-TPS-YOLO: Complete Implementation

🎯 Project Overview

This is the complete, production-ready implementation of the CBAM-STN-TPS-YOLO model described in the research paper "CBAM-STN-TPS-YOLO: Enhancing Agricultural Object Detection through Spatially Adaptive Attention Mechanisms".

🌱 Overview

CBAM-STN-TPS-YOLO integrates three key components:

• Spatial Transformer Networks (STN) for spatial invariance
• Thin-Plate Splines (TPS) for non-rigid deformation handling
• Convolutional Block Attention Module (CBAM) for feature attention

✅ What's Included

• ✅ Full Model Implementation (All 5 variants: YOLO, STN-YOLO, STN-TPS-YOLO, CBAM-STN-YOLO, CBAM-STN-TPS-YOLO)
• ✅ Complete Loss Functions (CIoU Loss, Distributed Focal Loss, Full YOLO Loss)
• ✅ Comprehensive Metrics (Precision, Recall, mAP, F1-Score with proper IoU calculation)
• ✅ Dataset Loading (PGP, MelonFlower, GlobalWheat with multi-spectral support)
• ✅ Data Augmentations (Rotation, Shear, Crop, Color Jitter with bbox transformation)
• ✅ Training Infrastructure (Multi-GPU support, early stopping, checkpointing)
• ✅ Evaluation Tools (Statistical analysis, confusion matrices, attention visualization)
• ✅ Inference Pipeline (Single image and batch prediction)
• ✅ Experimental Framework (Reproduces all paper results)
• ✅ Visualization Tools (TPS warping, attention maps, training curves)
• ✅ Edge Deployment (Optimized for Jetson platforms)

🚀 Quick Start

Installation

# Clone repository
git clone https://github.com/your-username/CBAM-STN-TPS-YOLO.git
cd CBAM-STN-TPS-YOLO

# Install dependencies
pip install -r requirements.txt

# Install package
pip install -e .

Reproduce Paper Results

# Run complete experimental suite (all models, all augmentations, 3 seeds each)
python experiments/run_experiments.py

# Statistical analysis
python experiments/statistical_analysis.py

# Single model training
python experiments/run_experiments.py --model CBAM-STN-TPS-YOLO --single

Quick Training

# Train best model
python -m src.training.trainer --config config/training_configs.yaml

Inference (using code)

# Single image
python -m src.inference.predict --checkpoint results/best_cbam_stn_tps_yolo.pth --input image.jpg

# Batch processing
python -m src.inference.predict --checkpoint results/best_cbam_stn_tps_yolo.pth --input images/ --output results/

📁 Complete Project Structure

CBAM-STN-TPS-YOLO/
├── README.md                          # This file
├── requirements.txt                   # Dependencies
├── setup.py                          # Package installation
├── .gitignore                        # Git ignore rules
├── LICENSE                           # MIT License
│
├── config/
│   ├── training_configs.yaml         # Training configurations
│   └── model_configs.yaml            # Model architecture configs
│
├── src/
│   ├── __init__.py
│   ├── models/
│   │   ├── __init__.py
│   │   ├── cbam.py                   # ✅ CBAM implementation
│   │   ├── stn_tps.py                # ✅ STN with TPS transformation
│   │   ├── yolo_backbone.py          # ✅ YOLO backbone with CBAM
│   │   ├── detection_head.py         # ✅ YOLO detection heads
│   │   └── cbam_stn_tps_yolo.py     # ✅ Complete model + variants
│   │
│   ├── data/
│   │   ├── __init__.py
│   │   ├── dataset.py                # ✅ PGP, MelonFlower, GlobalWheat datasets
│   │   ├── transforms.py             # ✅ Augmentations with bbox transforms
│   │   └── preprocessing.py          # ✅ Data preprocessing utilities
│   │
│   ├── training/
│   │   ├── __init__.py
│   │   ├── trainer.py                # ✅ Complete training infrastructure
│   │   ├── losses.py                 # ✅ CIoU, Focal, YOLO losses
│   │   └── metrics.py                # ✅ Detection metrics with NMS
│   │
│   ├── utils/
│   │   ├── __init__.py
│   │   ├── visualization.py          # ✅ Plotting, attention maps, TPS viz
│   │   └── evaluation.py             # ✅ Model evaluation tools
│   │
│   └── inference/
│       ├── __init__.py
│       └── predict.py                # ✅ Inference pipeline
│
├── experiments/
│   ├── run_experiments.py            # ✅ Complete experimental suite
│   ├── ablation_study.py             # ✅ Ablation experiments
│   └── statistical_analysis.py       # ✅ Statistical significance testing
│
├── data/                             # Dataset directory
│   ├── PGP/                          # Plant Growth & Phenotyping
│   ├── MelonFlower/                  # MelonFlower dataset
│   └── GlobalWheat/                  # GlobalWheat dataset
│
├── results/                          # Experimental results
│   ├── models/                       # Trained model checkpoints
│   ├── plots/                        # Generated visualizations
│   ├── experimental_results.json     # Complete results table
│   └── statistical_analysis.png      # Statistical plots
│
├── notebooks/                        # Analysis notebooks
│   ├── data_exploration.ipynb        # Dataset analysis
│   ├── model_analysis.ipynb          # Model behavior analysis
│   └── results_visualization.ipynb   # Results plotting
│
└── docs/                            # Documentation
    ├── installation.md               # Installation guide
    ├── usage.md                      # Usage examples
    ├── api_reference.md              # API documentation
    └── paper_reproduction.md         # Reproducing paper results

📊 Paper Results Reproduction

Main Results Table (Table in Paper)

Model	Accuracy	Precision	Recall	mAP	F1-Score	Inference Time
YOLO	84.86 ± 0.47	94.30 ± 0.56	89.21 ± 0.53	71.76 ± 1.03	91.68	16.25 ms
STN-YOLO	81.63 ± 1.53	95.34 ± 0.76	89.52 ± 0.57	72.56 ± 0.90	92.14	16.92 ms
STN-TPS-YOLO	82.48 ± 1.22	95.76 ± 0.81	89.70 ± 0.60	73.01 ± 0.88	92.41	15.18 ms
CBAM-STN-YOLO	82.73 ± 1.38	95.11 ± 0.73	89.89 ± 0.59	72.87 ± 0.81	92.46	14.69 ms
CBAM-STN-TPS-YOLO	83.24 ± 1.30	96.27 ± 0.72	90.28 ± 0.60	73.71 ± 0.85	92.78	14.22 ms

Key Improvements

• 12% reduction in false positives (improved precision)
• 1.9% improvement in mAP over baseline YOLO
• 13% faster inference compared to STN-YOLO
• Statistically significant improvements (p < 0.05) across all metrics

🧪 Experimental Features

1. Multi-Spectral Image Support

# Load PGP dataset with 4 spectral bands (580nm, 660nm, 730nm, 820nm)
dataset = PGPDataset(data_dir='data/PGP', multi_spectral=True)

2. TPS Visualization

# Visualize Thin-Plate Spline transformations
visualizer.visualize_tps_transformation(original_img, transformed_img)

3. Attention Map Analysis

# Visualize CBAM attention maps
visualizer.plot_attention_maps(image, channel_attention, spatial_attention)

4. Robustness Testing

# Test with different augmentations
test_augs = TestAugmentations()
transform = test_augs.get_transform('rotation_shear_crop')

🔬 Scientific Contributions

1. Novel TPS Integration

• Replaces rigid affine transformations with flexible Thin-Plate Splines
• Handles non-rigid deformations in plant structures
• Regularization parameter λ controls smoothness vs. flexibility

2. CBAM Enhancement

• Sequential channel and spatial attention
• Suppresses background noise effectively
• Lightweight design for edge deployment

3. Agricultural Optimization

• Multi-spectral image processing
• Occlusion-heavy dataset performance
• Real-time inference capability

📈 Performance Benchmarks

Inference Speed (NVIDIA Jetson Xavier)

• CBAM-STN-TPS-YOLO: 14.22 ms (70.4 FPS)
• STN-YOLO: 16.92 ms (59.1 FPS)
• YOLO Baseline: 16.25 ms (61.5 FPS)

Memory Usage

• Model Size: 45.2 MB
• Peak GPU Memory: 2.1 GB (training)
• Runtime Memory: 320 MB (inference)

Accuracy vs Speed Trade-off

• 13% faster than STN-YOLO
• 1.9% higher mAP than baseline
• 12% fewer false positives

💻 Hardware Requirements

Minimum Requirements

• GPU: NVIDIA GTX 1080 Ti (11GB VRAM) or equivalent
• RAM: 16GB system memory
• Storage: 100GB available space
• CPU: Intel i5-8400 / AMD Ryzen 5 2600 or better
• CUDA: Version 11.8 or higher

Recommended Requirements

• GPU: NVIDIA RTX 3090 (24GB VRAM) or RTX 4090
• RAM: 32GB system memory
• Storage: 500GB SSD
• CPU: Intel i7-10700K / AMD Ryzen 7 3700X or better
• CUDA: Version 12.1 or higher

Training Time Estimates

On RTX 3090 (24GB):

• YOLO baseline: ~8 hours (200 epochs)
• STN-YOLO: ~10 hours (200 epochs)
• STN-TPS-YOLO: ~14 hours (200 epochs)
• CBAM-STN-YOLO: ~12 hours (200 epochs)
• CBAM-STN-TPS-YOLO: ~16 hours (200 epochs)

On RTX 4090 (24GB):

• CBAM-STN-TPS-YOLO: ~12 hours (200 epochs)

Edge Deployment Compatibility

• NVIDIA Jetson Xavier NX: ✅ Supported (INT8 quantization recommended)
• NVIDIA Jetson AGX Orin: ✅ Fully supported
• Intel Neural Compute Stick: ⚠️ Limited support (ONNX export required)
• Google Coral TPU: ❌ Not supported (architecture incompatible)

🛠️ Development Features

1. Automatic Mixed Precision

# Enable AMP for faster training
config['mixed_precision'] = True

2. Multi-GPU Support

# Automatic multi-GPU detection
model = nn.DataParallel(model)

3. Experiment Tracking

# Wandb integration
config['use_wandb'] = True

4. Statistical Analysis

# Automatic significance testing
perform_statistical_analysis()

🔧 Customization Guide

Add New Dataset

class CustomDataset(Dataset):
    def __init__(self, data_dir, split='train'):
        # Implement dataset loading
        pass

    def __getitem__(self, idx):
        # Return image, targets, path
        pass

Modify Model Architecture

# Create custom model variant
model = create_model(
    model_type='CBAM-STN-TPS-YOLO',
    num_classes=5,  # Custom number of classes
    num_control_points=30,  # More TPS control points
    backbone_channels=[64, 128, 256, 512, 1024]  # Larger backbone
)

Custom Loss Function

class CustomLoss(nn.Module):
    def __init__(self):
        super().__init__()
        # Implement custom loss

    def forward(self, predictions, targets):
        # Calculate custom loss
        pass

📚 API Reference

Model Creation

from src.models import create_model, CBAM_STN_TPS_YOLO

# Create specific model variant
model = create_model(
    model_type='CBAM-STN-TPS-YOLO',
    num_classes=5,
    input_channels=3,
    num_control_points=20,
    backbone_type='darknet53'
)

# Direct model instantiation
model = CBAM_STN_TPS_YOLO(
    num_classes=5,
    num_control_points=20,
    cbam_reduction_ratio=16,
    tps_regularization=0.1
)

Dataset Loading

from src.data import create_agricultural_dataloader, PGPDataset

# Create data loader
train_loader = create_agricultural_dataloader(
    data_dir='data/PGP',
    split='train',
    batch_size=16,
    image_size=640,
    augmentation_type='advanced'
)

# Direct dataset usage
dataset = PGPDataset(
    data_dir='data/PGP',
    split='train',
    multi_spectral=True,
    transform=transforms
)

Training

from src.training import CBAMSTNTPSYOLOTrainer

# Initialize trainer
trainer = CBAMSTNTPSYOLOTrainer(config, model_type='CBAM-STN-TPS-YOLO')

# Train model
best_mAP = trainer.train()

# Resume training
trainer.resume_from_checkpoint('path/to/checkpoint.pth')

Inference

from src.inference import ModelPredictor

# Initialize predictor
predictor = ModelPredictor(
    model_path='path/to/model.pth',
    device='cuda',
    confidence_threshold=0.5
)

# Single image prediction
results = predictor.predict_image('path/to/image.jpg')

# Batch prediction
results = predictor.predict_batch('path/to/images/', 'path/to/output/')

Visualization

from src.utils.visualization import Visualizer

# Initialize visualizer
viz = Visualizer(class_names=['Cotton', 'Rice', 'Corn'])

# Plot attention maps
viz.plot_attention_maps(image, attention_weights)

# Visualize TPS transformation
viz.visualize_tps_transformation(original_img, transformed_img, control_points)

# Plot training curves
viz.plot_training_curves(train_losses, val_losses, metrics)

📊 Results

Our model achieves the following improvements over baseline YOLO:

• Precision: 96.27% (+2.0%)
• Recall: 90.28% (+1.1%)
• mAP: 73.71% (+1.9%)
• Inference Time: 14.22ms (13% faster)

🏗️ Architecture

Input Image (Multi-spectral)
    ↓
STN with TPS Transformation
    ↓
CBAM Attention (Channel + Spatial)
    ↓
YOLO Backbone + Detection Head
    ↓
Bounding Boxes + Classes

📈 Key Features

• Multi-spectral Image Support: Handles 4-band spectral imaging (580nm, 660nm, 730nm, 820nm)
• Pseudo-RGB Generation: Converts multi-spectral to RGB for pre-trained model compatibility
• Robust Augmentation Testing: Evaluates performance under rotation, shear, and crop transformations
• Edge Deployment Ready: Optimized for NVIDIA Jetson platforms
• Comprehensive Evaluation: Statistical significance testing across multiple runs

🎯 Applications

• Plant phenotyping and growth monitoring
• Crop disease detection
• Precision agriculture automation
• Smart farming systems
• Automated greenhouse monitoring

📖 Citation

If you use this code in your research, please cite our paper:

@misc{praveen2025cbamstntpsyoloenhancingagriculturalobject,
      title={CBAM-STN-TPS-YOLO: Enhancing Agricultural Object Detection through Spatially Adaptive Attention Mechanisms},
      author={Satvik Praveen and Yoonsung Jung},
      year={2025},
      eprint={2506.07357},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2506.07357},
}

🛠️ Troubleshooting

Common Issues and Solutions

1. GPU Memory Issues

Problem: CUDA out of memory error during training Solutions:

# Reduce batch size
python experiments/run_experiments.py --batch_size 8

# Enable gradient checkpointing
python experiments/run_experiments.py --gradient_checkpointing

# Use mixed precision training
python experiments/run_experiments.py --mixed_precision

2. Installation Issues

Problem: PyTorch installation fails or CUDA version mismatch Solutions:

# Check CUDA version
nvidia-smi

# Install specific PyTorch version for CUDA 11.8
pip install torch==1.13.1+cu118 torchvision==0.14.1+cu118 -f https://download.pytorch.org/whl/torch_stable.html

# For CUDA 12.1
pip install torch==2.0.1+cu121 torchvision==0.15.2+cu121 -f https://download.pytorch.org/whl/torch_stable.html

3. Dataset Loading Issues

Problem: Dataset not found or incorrect format Solutions:

# Verify dataset structure
python -c "from src.data import verify_dataset_structure; verify_dataset_structure('data/PGP')"

# Download datasets automatically
python scripts/download_datasets.py --dataset all

# Validate dataset annotations
python scripts/validate_annotations.py --data_dir data/PGP

4. Training Instability

Problem: Loss not converging or NaN values Solutions:

# Reduce learning rate
config['learning_rate'] = 0.0005

# Increase warmup epochs
config['warmup_epochs'] = 10

# Check data preprocessing
config['verify_data'] = True

5. Poor Performance

Problem: Model performance below expected results Solutions:

# Verify data augmentation
python experiments/test_augmentations.py

# Check model configuration
python experiments/verify_model_config.py

# Run ablation study
python experiments/ablation_study.py --quick

Performance Optimization Tips

Memory Optimization

• Use gradient accumulation for larger effective batch sizes
• Enable memory-efficient attention mechanisms
• Use checkpoint saving to resume interrupted training

Speed Optimization

• Use DataLoader with multiple workers (num_workers=4-8)
• Enable pin_memory for faster GPU transfer
• Use mixed precision training (AMP)

Model Optimization

• Experiment with different TPS control point numbers (10-30)
• Adjust CBAM reduction ratios (8, 16, 32)
• Try different backbone architectures

Getting Help

If you encounter issues not covered here:

1. Check the Issues page
2. Search existing discussions
3. Create a new issue with:
   • Error message and full traceback
   • System information (python --version, nvidia-smi)
   • Minimal code to reproduce the issue
   • Configuration file used

5. Environment Setup Section (add after Installation)

🔧 Environment Setup

Method 1: Conda Environment (Recommended)

# Create conda environment
conda create -n cbam-stn-tps-yolo python=3.9
conda activate cbam-stn-tps-yolo

# Install PyTorch with CUDA support
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

# Install other requirements
pip install -r requirements.txt

# Install package in development mode
pip install -e .

Method 2: Virtual Environment

# Create virtual environment
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Upgrade pip
pip install --upgrade pip

# Install requirements
pip install -r requirements.txt

# Install package
pip install -e .

Method 3: Docker (Production)

# Build Docker image
docker build -t cbam-stn-tps-yolo .

# Run with GPU support
docker run --gpus all -it cbam-stn-tps-yolo

# Mount data directory
docker run --gpus all -v /path/to/data:/app/data -it cbam-stn-tps-yolo

Method 4: Google Colab

# In Colab notebook
!git clone https://github.com/your-username/CBAM-STN-TPS-YOLO.git
%cd CBAM-STN-TPS-YOLO
!pip install -r requirements.txt
!pip install -e .

# Check GPU availability
import torch
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'None'}")

Verification

# Test installation
python -c "
import torch
import src
from src.models import create_model
print('✅ Installation successful!')
print(f'PyTorch version: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
model = create_model('CBAM-STN-TPS-YOLO', num_classes=5)
print(f'✅ Model creation successful!')
"

🤝 Contributing

We welcome contributions! Please see CONTRIBUTING.md for guidelines.

Development Setup

# Install development dependencies
pip install -e ".[dev]"

# Run tests
python -m pytest tests/

# Code formatting
black src/
isort src/

📄 License

This project is licensed under the MIT License - see LICENSE for details.

🙏 Acknowledgments

• Texas A&M AgriLife for support
• Texas A&M High Performance Research Computing (HPRC) for computational resources
• Zambre et al. for the original STN-YOLO implementation

📞 Contact

• Satvik Praveen - Email
• Yoonsung Jung - Email

---

🎯 Next Steps

1. Download datasets and place in data/ directory
2. Run experiments to reproduce paper results
3. Explore notebooks for detailed analysis
4. Customize models for your specific use case
5. Deploy to edge devices using provided optimization tools

Ready to revolutionize agricultural object detection! 🌱🚀

---

⭐ If you find this work useful, please star this repository!

---

Authors: Satvik Praveen, Yoonsung Jung
Institution: Texas A&M University