This repository contains the official implementation of our paper:

Restora-Flow: Mask-Guided Image Restoration with Flow Matching
Accepted at WACV 2026.

# Clone the repository git clone <repo-url> cd Restora-Flow # Create a new conda environment conda create -n restora-flow python=3.10 -y # Activate the conda environment conda activate restora-flow # Install PyTorch pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 # Install required packages pip install -r requirements.txt

Supported datasets: CelebA, AFHQ-Cat, COCO

• Download: https://www.kaggle.com/datasets/jessicali9530/celeba-dataset
• Extract all contents into:

  natural/data/celeba 

• Download: https://www.kaggle.com/datasets/dimensi0n/afhq-512/data
• Place test images into:

  natural/data/afhq_cat/test/cat 

• Download: http://images.cocodataset.org/zips/val2017.zip
• Place images into:

  natural/data/coco/val 

natural/data contains .txt files listing the image IDs used for evaluation in the paper.

Download all models:

cd natural chmod +x download_all.sh ./download_all.sh

Or download individually:

cd natural bash download.sh celeba-ot bash download.sh celeba-ddpm bash download.sh afhq-cat-ot bash download.sh afhq-cat-ddpm bash download.sh coco-ot bash download.sh coco-ddpm

Checkpoints are stored in: natural/model_checkpoints/{dataset}/gaussian/{model_type}

Before running experiments, set the repository base path in natural/src/dataloaders.py:

base_path = "<absolute-path-to-Restora-Flow>"

Supported:

• Model types: ot, ddpm
• Problems: denoising, box_inpainting, superresolution, random_inpainting
• Methods: repaint, ddnm, ot_ode, flow_priors, d_flow, pnp_flow, restora_flow

cd natural python main.py --opts dataset celeba eval_split test model ot problem box_inpainting method restora_flow ode_steps 64 correction_steps 1 max_batch 1 batch_size_ip 1

natural/script_test.sh 

Results are saved to:

natural/results/{dataset}/{model_type}/{problem}/ 

Supported datasets: X-ray Hand

1. Download the X-ray hand dataset from the original source:
   https://www.ipilab.org/Research/BAA/BAAindex.html

2. Place images into:

medical/dataset/xray_hand/images 

3. (Optional) Adjust the image extension in:

medical/dataset.py 

cd medical chmod +x download.sh bash download.sh xray-hand-flow bash download.sh xray-hand-ddpm

Checkpoints are saved to: medical/model_checkpoints/xray_hand/{model_type}/full

Supported:

• Model types: flow, ddpm
• Problems: denoising, box_inpainting, superresolution, occlusion_removal
• Methods: RePaint, DDNM, OT-ODE, Flow-Priors, D-Flow, PnP-Flow, Restora-Flow

cd medical python main.py --model_type {flow|ddpm} --method {method} --problem {problem}

Example:

python main.py --model_type flow --method Restora-Flow --problem box_inpainting

Configuration options: medical/configs/restoration_config.yaml

Output is saved to: medical/exports/hand_samples/{model_type}/{problem}/{method}/{timestamp}/

@article{hadzic2025restoraflow, title={Restora-Flow: Mask-Guided Image Restoration with Flow Matching}, author={Hadzic, Arnela and Thaler, Franz and Bogensperger, Lea and Joham, Simon Johannes and Urschler, Martin}, journal={arXiv preprint arXiv:2511.20152}, year={2025} } 

MedicalDataAugmentationTool taken from:
https://github.com/christianpayer/MedicalDataAugmentationTool

GaussianDiffusion implementation based on:
https://github.com/mobaidoctor/med-ddpm

Natural image framework builds upon:
https://github.com/annegnx/PnP-Flow

CelebA and AFHQ-Cat flow models are from PnP-Flow.
Other pretrained models were trained as part of this project.