File size: 9,824 Bytes
34c53c1 2ad4160 34c53c1 2ad4160 7f901ee 2ad4160 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | ---
library_name: pytorch
pipeline_tag: image-to-image
license: cc-by-nc-sa-4.0
language:
- en
tags:
- image-inpainting
- medfe
- partial-convolution
- computer-vision
---
<div align="center">
# MEDFE
### Mutual Encoder–Decoder with Feature Equalizations
**ECCV 2020 · Modern PyTorch implementation**
[Paper](https://www.ecva.net/papers/eccv_2020/papers_ECCV/papers/123470715.pdf) · [Weights](https://huggingface.co/KumaPower/MEDFE) · [Quick start](#quick-start) · [Training](#training) · [Datasets](https://github.com/KumapowerLIU/Rethinking-Inpainting-MEDFE/blob/master/docs/data.md) · [Architecture](#architecture)
Hongyu Liu · Bin Jiang · Yibing Song · Wei Huang · Chao Yang
</div>

*Selected MEDFE examples. Each row shows masked input → completion → ground truth. [Checkpoint details](https://github.com/KumapowerLIU/Rethinking-Inpainting-MEDFE/blob/master/docs/checkpoint.md).*
This repository is a new PyTorch implementation of **MEDFE**, which jointly recovers image structure and texture through a mutual encoder–decoder and feature equalization. It provides batched training, reproducible random masks, explicit configuration, resumable checkpoints, and offline inference.
The network combines a **PartialConv residual encoder** with a **GroupNorm/SiLU residual decoder**, using nearest-neighbor upsampling and ordinary convolutions. Structure and texture branches fill missing features through multi-scale partial convolutions, followed by channel equalization and bilateral propagation. A single forward pass produces the completed RGB image.
## Quick start
Python **3.10+** is required. The training environment used PyTorch **2.6.0**, torchvision **0.21.0** and CUDA **12.4**. For Linux/NVIDIA:
```bash
git clone https://github.com/KumapowerLIU/Rethinking-Inpainting-MEDFE.git
cd Rethinking-Inpainting-MEDFE
python3 -m venv .venv
source .venv/bin/activate
python -m pip install torch==2.6.0 torchvision==0.21.0 \
--index-url https://download.pytorch.org/whl/cu124
python -m pip install -e .
medfe --help
```
For macOS or CPU, install the matching [PyTorch wheels](https://pytorch.org/get-started/previous-versions/#v260) before installing this package.
Download the **step 120,994** generator from [Hugging Face](https://huggingface.co/KumaPower/MEDFE), then run:
```bash
mkdir -p checkpoints
curl -L --fail \
https://huggingface.co/KumaPower/MEDFE/resolve/main/medfe-sd-partial-step120994.pt \
-o checkpoints/medfe-sd-partial.pt
medfe infer --checkpoint checkpoints/medfe-sd-partial.pt \
--image examples/masked.png --mask examples/mask.png \
--output outputs/completed.png --device cuda:0
```
**White mask pixels are holes; black pixels are known.** Image and mask must have matching dimensions. Inference preserves the original size through padding/cropping, copies known pixels exactly, and writes a lossless PNG. `--device cpu` is supported. The download is a **579.6 MB FP32 inference bundle**; [checkpoint details](https://github.com/KumapowerLIU/Rethinking-Inpainting-MEDFE/blob/master/docs/checkpoint.md) include its checksum and export instructions.
## Architecture

*Architecture overview, with a layout inspired by [Figure 2 of the paper](https://www.ecva.net/papers/eccv_2020/papers_ECCV/papers/123470715.pdf#page=5). Scale labels are spatial side lengths for a 256 × 256 input. The image tiles are schematic; model results are shown above.*
The generator is trained jointly. Equalized features are projected to each encoder scale and added to the encoder features; the resulting feature pyramid supplies the decoder input and skip connections. See the [architecture notes](https://github.com/KumapowerLIU/Rethinking-Inpainting-MEDFE/blob/master/docs/sd-partial-backbone.md) for feature scales, mask propagation and normalization.
## Training
The current recipe is [`configs/sd-partial-csa-mix-large-batch.yaml`](https://github.com/KumapowerLIU/Rethinking-Inpainting-MEDFE/blob/master/configs/sd-partial-csa-mix-large-batch.yaml).
| Setting | Current experiment |
| --- | --- |
| GPU | **1 × NVIDIA H100 80 GB HBM3** |
| Image size / batch | **256 × 256 / 128**, one GPU |
| Precision | BF16 training; FP32 perceptual computation and validation |
| Memory | About **33 GiB** process usage observed; **38 GiB** allocator budget |
| Measured short-test peak | 25.925 GiB allocated / 30.551 GiB reserved; synthetic full D/G test |
| Optimizer | Adam, lr `0.0002`, betas `(0.5, 0.999)` |
| Masks | Rectangles, brush strokes, or mixtures; 10%–40% missing area |
| Stop point | **120,994 steps**, paused with a complete training checkpoint |
The 80 GB card capacity is not the measured requirement. Memory varies with inputs, backends and image size; reduce batch size for smaller GPUs. [Training details](https://github.com/KumapowerLIU/Rethinking-Inpainting-MEDFE/blob/master/docs/training.md) include the loss definitions, warmup, resume rules and hardware measurements.
### 1. Prepare your images
Use separate training and validation folders. The indexer reads existing files; it does not download or extract a dataset.
```bash
python scripts/build_image_index.py --images-root /path/to/train \
--manifest data/train.jsonl
python scripts/build_image_index.py --images-root /path/to/validation \
--manifest data/validation.jsonl
```
The indexed reader resizes RGB to 256², computes a true RTV structure target, and caches it on first use. See [data preparation](https://github.com/KumapowerLIU/Rethinking-Inpainting-MEDFE/blob/master/docs/data.md) for storage requirements, archive formats, face sampling and the alternative paired-data protocol.
### 2. Prepare LPIPS weights
```bash
export TORCH_HOME="$PWD/weights/torch"
medfe prepare-weights --lpips-only --output weights
```
This explicitly fetches missing VGG-LPIPS assets. The current discriminators are randomly initialized PatchGANs and need no pretrained D. Training itself defaults to **no implicit downloads**.
### 3. Train and resume
```bash
medfe train --config configs/sd-partial-csa-mix-large-batch.yaml \
--set train.device=cuda:0 --max-steps 20
medfe train --config configs/sd-partial-csa-mix-large-batch.yaml \
--set train.device=cuda:0 \
--resume runs/medfe-sd-partial-csa-mix-large-batch/latest.pt
```
The first command is a short startup check; `--max-steps` is an **absolute** step limit. The second continues the same run. Checkpoints contain both discriminators, optimizer states, perceptual weights, RNG state and the data cursor. SIGINT/SIGTERM finish the current update and save a complete checkpoint. Source/config/data checks protect resume compatibility; archived runs use their frozen source for exact continuation.
| Generator loss | Weight |
| --- | ---: |
| Full-image raw-prediction L1 | 1 |
| Structure auxiliary L1 | 1 |
| Texture auxiliary L1 | 1 |
| VGG-LPIPS on the completed image | 1 |
| Global + local RaLSGAN | Ramps to **0.002** |
GAN is disabled for the first 5,000 updates, then ramps linearly over 10,000 updates. Separate VGG perceptual/style losses are disabled in this recipe.
## Data used for this checkpoint
The run uses the same source mix and split as the modern CSA experiment:
| Source | Training entries | Obtain the source data |
| --- | ---: | --- |
| Unsplash-derived photographs | 296,145 | [Unsplash Dataset](https://unsplash.com/data) · [access/schema](https://github.com/unsplash/datasets) |
| Places365-Standard | 901,730, half of the training archive | [CSAIL dataset repository](https://github.com/CSAILVision/places365) |
| HDTF talking-head frames | 299,469 entries from 30,128 unique frames | [HDTF dataset and preprocessing](https://github.com/MRzzm/HDTF) |
**1,497,344 sampling entries; 1,228,003 distinct source records.** Face frames repeat 9 or 10 times for approximately 20% sampling weight. Validation uses 2,000 held-out Unsplash records. The training masks are generated online. The local selection/manifests are not a public benchmark split; source access and reproducibility details are in [datasets](https://github.com/KumapowerLIU/Rethinking-Inpainting-MEDFE/blob/master/docs/data.md).
## Repository guide
| Location | Purpose |
| --- | --- |
| `src/medfe/model.py`, `blocks.py`, `sd_backbone.py` | Generator, feature equalization, PartialConv and decoder |
| `src/medfe/trainer.py`, `losses.py`, `discriminators.py` | Joint training, loss definitions and PatchGANs |
| `src/medfe/indexed_data.py`, `rtv.py`, `masks.py` | Data decoding, cached structure targets and random masks |
| `src/medfe/evaluation.py`, `inference.py` | Image-preserving inference, evaluation and portable export |
| `configs/` | Training configurations |
| `legacy/` | Archived source files |
```bash
python -m pip install -e '.[dev]'
ruff check src tests scripts
python -m pytest tests -q
```
Large holes can still lead to repeated textures, implausible faces and missing body parts. See the [validation records](https://github.com/KumapowerLIU/Rethinking-Inpainting-MEDFE/blob/master/docs/validation.md) for implementation checks and image-quality observations.
## Citation and license
```bibtex
@inproceedings{liu2020rethinking,
title={Rethinking Image Inpainting via a Mutual Encoder-Decoder with Feature Equalizations},
author={Liu, Hongyu and Jiang, Bin and Song, Yibing and Huang, Wei and Yang, Chao},
booktitle={European Conference on Computer Vision},
year={2020}
}
```
[CC BY-NC-SA 4.0](https://github.com/KumapowerLIU/Rethinking-Inpainting-MEDFE/blob/master/LICENSE.md). Dataset and pretrained-weight terms remain those of their respective providers.
|