github 項目推薦：用 edge-connect 進(jìn)行圖像修復(fù)

本文作者： skura

2019-02-20 19:12

導(dǎo)語：線條優(yōu)先，顏色次之，這是一種新的圖像修復(fù)方法，可以很好的還原圖像。

雷鋒網(wǎng) AI 科技評論按，用對抗性邊緣學(xué)習(xí)修復(fù)生成圖像是一種新的圖像修復(fù)方法，它可以更好地復(fù)制填充區(qū)域，它的細(xì)節(jié)部分展現(xiàn)了開發(fā)者對藝術(shù)工作者工作方式的理解：線條優(yōu)先，顏色次之。對應(yīng)的論文在 arxiv 上可以查看：https://arxiv.org/abs/1901.00212。

文中提出了一種 2 階對抗式邊緣連接模型，該模型由一個邊緣生成器和一個圖像完成網(wǎng)絡(luò)組成。邊緣生成器先描繪出圖像缺失區(qū)域（規(guī)則和不規(guī)則）的邊緣，圖像完成網(wǎng)絡(luò)先驗使用描繪出的邊緣填充缺失區(qū)域。論文對該系統(tǒng)進(jìn)行了詳細(xì)的描述。

github 項目推薦：用 edge-connect 進(jìn)行圖像修復(fù)

（a）輸入有缺失區(qū)域的圖像，缺失區(qū)域用白色表示。（b）計算得到的邊緣，使用 Canny 邊緣檢測器計算（針對可用區(qū)域）黑色繪制的邊緣；而藍(lán)色顯示的邊緣由邊緣生成器網(wǎng)絡(luò)描繪。（c）擬用方法的圖像修復(fù)結(jié)果。

準(zhǔn)備：

Python 3

PyTorch 1.0

NVIDIA GPU + CUDA cuDNN

安裝：

復(fù)制下面這個 repo：

git clone https://github.com/knazeri/edge-connect.gitcd edge-connect
cd edge-connect

從 http://pytorch.org 安裝 PyTorch 及其相關(guān)依賴。

安裝 python 文件：

pip install -r requirements.txt

數(shù)據(jù)集

1. 圖像

這里使用 Places2, CelebA 以及 Paris Street-View 數(shù)據(jù)集。從官網(wǎng)下載數(shù)據(jù)集，在整個數(shù)據(jù)集上訓(xùn)練模型。

下載完成后，運行 scripts/flist.py 這個文件來生成訓(xùn)練、測試和驗證集文件列表。例如，要在 Places2 數(shù)據(jù)集上生成訓(xùn)練集文件列表，請運行：

mkdir datasets
python ./scripts/flist.py --path path_to_places2_train_set --output ./datasets/places_train.flist

2 .不規(guī)則掩膜

我們的模型是在 Liu 等人提供的不規(guī)則掩模數(shù)據(jù)集上進(jìn)行訓(xùn)練的，你可以從他們的網(wǎng)站上下載公開的不規(guī)則掩膜數(shù)據(jù)集。雷鋒網(wǎng)

請使用 scripts/flist.py 生成上述訓(xùn)練、測試和驗證集掩膜文件列表。

開始

使用以下鏈接下載預(yù)先訓(xùn)練的模型，并將其復(fù)制到./checkpoints 目錄下。

Places2 | CelebA | Paris-StreetView

或者，你可以運行以下腳本自動下載預(yù)訓(xùn)練的模型：

bash ./scripts/download_model.sh

1 .訓(xùn)練

要訓(xùn)練模型，請創(chuàng)建一個類似于示例配置文件的 config.yaml 文件，并將其復(fù)制到檢查點目錄下。有關(guān)模型配置的詳細(xì)信息，請參閱配置指南。

EdgeConnect 的訓(xùn)練分為三個階段：1）邊緣模型的訓(xùn)練；2）內(nèi)部模型的訓(xùn)練；3）聯(lián)合模型的訓(xùn)練。訓(xùn)練模型：

python train.py --model [stage] --checkpoints [path to checkpoints]

例如，要在./checkpoints/places2 目錄下的 places2 數(shù)據(jù)集上訓(xùn)練邊緣模型：

python train.py --model 1 --checkpoints ./checkpoints/places2

模型的收斂性因數(shù)據(jù)集而異。例如，Places2 數(shù)據(jù)集在兩個時期中的一個就能聚合，而較小的數(shù)據(jù)集（如 CelebA）則需要將近 40 個時期才能聚合。你可以通過更改配置文件中的 MAX_ITERS 值來設(shè)置訓(xùn)練迭代次數(shù)。雷鋒網(wǎng)

2 .測試

要測試模型，請創(chuàng)建一個與示例配置文件類似的 config.yaml 文件，并將其復(fù)制到檢查點目錄下。

你可以在所有三個階段上測試模型：邊緣模型、內(nèi)部模型和聯(lián)合模型。在每種情況下，都需要提供一個輸入圖像（帶掩膜的圖像）和一個灰度掩膜文件。請確保掩膜文件覆蓋輸入圖像中的整個掩膜區(qū)域。測試模型：

python test.py \
--model [stage]
--checkpoints [path to checkpoints] \
--input [path to input directory or file] \
--mask [path to masks directory or mask file] \
--output [path to the output directory]

我們在./examples 目錄下提供了一些測試示例，請下載預(yù)訓(xùn)練模型并運行：

python test.py \
--checkpoints ./checkpoints/places2
--input ./examples/places2/images
--mask ./examples/places2/masks
--output ./checkpoints/results

此腳本將在./examples/places2/images 中使用和./examples/places2/mask 對應(yīng)的掩膜圖像，并將結(jié)果保存在./checkpoints/results 目錄中。默認(rèn)情況下，test.py 腳本在階段 3 上運行（--model=3）。

3 .評估

要評估模型，你需要首先在測試模式下對 validation 集運行模型，并將結(jié)果保存到磁盤上。我們提供了一個實用程序./scripts/metrics.py，使用 PSNR, SSIM 和 Mean Absolute Error 評估模型：