ISTD Shadow Removal Model Training

Deep learning shadow removal model training code based on ISTD dataset.

Dataset Description

ISTD dataset contains:

• train_A/: Original shadowed images (RGB, 640x480)
• train_B/: Shadow masks (grayscale, values 0 and 255)
• train_C/: Shadow-free target images (RGB, 640x480)
• test_A/, test_B/, test_C/: Test set with same structure

Total 5,610 images: 3,990 training samples and 1,620 test samples.

Installation

pip install -r requirements.txt

Usage

Basic Training

python train.py

Custom Parameter Training

python train.py \
    --epochs 50 \
    --batch_size 16 \
    --learning_rate 1e-4 \
    --model_type unet \
    --save_dir ./my_checkpoints

Evaluation Only

python train.py --eval_only --resume ./checkpoints/model_best.pth

Resume Training from Checkpoint

python train.py --resume ./checkpoints/checkpoint_epoch_20.pth

Main Parameters

Parameter	Default	Description
--epochs	100	Number of training epochs
--batch_size	8	Batch size
--learning_rate	1e-4	Learning rate
--model_type	unet	Model type (unet/resnet)
--image_size	480 640	Image size [height width]
--save_every	10	Save model every N epochs
--device	auto	Training device (auto/cpu/cuda)

Model Architecture

U-Net Model

• Based on classic U-Net architecture
• 4-channel input: RGB image(3) + shadow mask(1)
• 3-channel output: shadow-free RGB image
• Supports bilinear interpolation upsampling

ResNet Model

• Based on ResNet encoder-decoder architecture
• Uses residual blocks for improved training stability
• Transpose convolution upsampling

Loss Function

Composite loss function includes:

• L1 Loss: Pixel-level reconstruction loss
• Perceptual Loss: High-level semantic loss based on VGG features
• Shadow Region Weighting: Higher weight for shadow regions

Output Files

After training, the following files will be generated in the specified directory:

checkpoints/
├── model_best.pth          # Best model
├── checkpoint_latest.pth   # Latest checkpoint
└── checkpoint_epoch_X.pth  # Epoch checkpoints

results/
├── training_history.png    # Training history curves
├── final_metrics.txt      # Final evaluation metrics
└── sample_X.png           # Sample result visualizations

Evaluation Metrics

• MSE: Mean Squared Error
• PSNR: Peak Signal-to-Noise Ratio
• SSIM: Structural Similarity Index
• Shadow_MSE/PSNR: Shadow region metrics
• NonShadow_MSE/PSNR: Non-shadow region metrics

Visualization Results

The training script automatically generates visualization results including:

1. Shadowed input image
2. Shadow mask
3. Shadow region annotation
4. Target image (ground truth)
5. Prediction result
6. Difference map

Hardware Requirements

• GPU: Recommended CUDA-compatible GPU
• Memory: At least 8GB RAM
• Storage: At least 5GB available space

Training Tips

1. Batch Size: Adjust based on GPU memory, recommend 8-16
2. Learning Rate: Start with 1e-4, adjust based on training progress
3. Data Augmentation: Can add random cropping, rotation, etc.
4. Early Stopping: Monitor validation loss to avoid overfitting

Troubleshooting

Out of Memory

• Reduce batch_size
• Reduce image_size
• Use CPU training

Slow Training Speed

• Increase num_workers
• Use GPU training
• Reduce image size

Model Not Converging

• Adjust learning rate
• Check data preprocessing
• Try different model architectures

File Structure

ISTD/
├── train.py              # Main training script
├── data_loader.py        # Data loader
├── model.py             # Model definition
├── trainer.py           # Trainer
├── utils.py             # Utility functions
├── requirements.txt     # Dependencies
├── README.md           # Documentation
├── train/              # Training data
│   ├── train_A/        # Shadowed images
│   ├── train_B/        # Shadow masks
│   └── train_C/        # Target images
└── test/               # Test data
    ├── test_A/
    ├── test_B/
    └── test_C/