TransWeather

repository·main·Indexed 42 days ago

https://github.com/jeya-maria-jose/transweather

A transformer-based model for the all-in-one restoration of images degraded by adverse weather conditions such as rain, fog, and snow. It utilizes a single encoder-decoder architecture to handle multiple weather types efficiently and supports fine-tuning for individual restoration tasks via the train-individual.py script.

Tokens
754
Snippets
3
Records
3
Agent score
3%

What's inside TransWeather

  1. Organize datasets for TransWeather

    main

    TransWeather requires datasets to be organized in a specific directory structure for training and testing. Each dataset should be placed within a train or test directory under data/, containing subdirectories for input (degraded images) and gt (ground truth clean images), along with a .txt file listing the filenames.

    Required Directory Structure:

    TransWeather
    ├── data 
    │   ├── train # Training  
    │   │   ├── <dataset_name>   
    │   │   │   ├── input         # degraded images (e.g., rain images) 
    │   │   │   └── gt            # clean images
    │   │   └── dataset_filename.txt
    │   └── test  # Testing         
    │   │   ├── <dataset_name>          
    │   │   │   ├── input         # degraded images
    │   │   │   └── gt            # clean images
    │   │   └── dataset_filename.txt
    TransWeather
    ├── data 
    │   ├── train # Training  
    │   │   ├── <dataset_name>   
    │   │   │   ├── input         # rain images 
    │   │   │   └── gt            # clean images
    │   │   └── dataset_filename.txt
    │   └── test  # Testing         
    │   │   ├── <dataset_name>          
    │   │   │   ├── input         # rain images 
    │   │   │   └── gt            # clean images
    │   │   └── dataset_filename.txt
  2. Install TransWeather

    main

    To use TransWeather, ensure you are using Python 3.6.13 and CUDA >= 10.1. You can install the project by cloning the repository and using Conda or Pip.

    git clone https://github.com/jeya-maria-jose/TransWeather
    cd TransWeather
    conda env create -f environment.yml
    conda activate transweather

    Using Pip

    If you prefer Pip, install the following specific versions to ensure compatibility:

    • timm==0.3.2
    • mmcv-full==1.2.7
    • torch==1.7.1
    • torchvision==0.8.2
    • opencv-python==4.5.1.48
  3. Fine-tune TransWeather for individual weather tasks

    main

    While TransWeather is designed as an all-in-one model, it can be modified to focus on a single restoration task (e.g., just rain removal) by removing the transformer decoder. To train the TransWeather-encoder on a specific dataset (like Rain-800), organize the data according to the standard dataset format and use the train-individual.py script.

    Note: The network used in this mode is a sub-section of the original architecture that excludes weather queries.

    python train-individual.py -train_batch_size 32 -exp_name Transweather-finetune -epoch_start 0 -num_epochs 250