cdotad / alphagan-matting Goto Github PK

This project is an unofficial implementation of AlphaGAN: Generative adversarial networks for natural image matting published at the BMVC 2018

Python 89.24% Jupyter Notebook 10.76%

pytorch gan matting

alphagan-matting's Introduction

AlphaGAN

This project is an unofficial implementation of AlphaGAN: Generative adversarial networks for natural image matting published at the BMVC 2018. As for now, the result of my experiment is not as good as the paper's.

Dataset

Adobe Deep Image Matting Dataset

Follow the instruction to contact the author for the dataset

You might need to follow the method mentioned in the Deep Image Matting to generate the trimap using the alpha mat.

The trimap are generated while the data are loaded.

import numpy as np
import cv2 as cv

def generate_trimap(alpha):
   k_size = random.choice(range(2, 5))
   iterations = np.random.randint(5, 15)
   kernel = cv.getStructuringElement(cv.MORPH_ELLIPSE, (k_size, k_size))
   dilated = cv.dilate(alpha, kernel, iterations=iterations)
   eroded = cv.erode(alpha, kernel, iterations=iterations)
   trimap = np.zeros(alpha.shape, dtype=np.uint8)
   trimap.fill(128)

   trimap[eroded >= 255] = 255
   trimap[dilated <= 0] = 0

   return trimap

See scripts/MattingTrain.ipynb and scripts/MattingTest.ipynb to compose the training/testing set.

The Dataset structure in my project

Train
  ├── alpha  # the alpha ground-truth
  ├── fg     # the foreground image
  ├── merged_cv  # the real image composed by the fg & bg
MSCOCO
  ├── train2014 # the background image

Running the Codes

   python train.py --dataroot ${YOUR_DIM_DATASET_ROOT} \
                     --training_file ${THE TRAINING FILE OF THE DIM DATASET}

Differences from the original paper

SyncBatchNorm instead of pytorch original BatchNorm when use multi GPU.
Training batch_size = 1 [1] [2]
Using GroupNorm [2]
Using Warmup [3] [4]

Records

4 GPUS 32 batch size, and SyncBatchNorm

Achieved SAD=78.22 after 21 epoches.

1 GPU 1 batch size, and GroupNorm

Achieved SAD=68.61 MSE=0.03189 after 33 epoches.
Achieved SAD=61.9 MSE=0.02878 after xx epoches.

Results

image	trimap	alpha(predicted)

Acknowledgments

My code is inspired by:

[1] pytorch-deep-image-matting
[2] FBA-Matting
[3] GCA-Matting
[4] reid-strong-baseline
pytorch-CycleGAN-and-pix2pix
pytorch-book chapter7 generate anime head portrait with GAN
pytorch-deeplab-xception
indexnet_matting

alphagan-matting's People

Contributors

Stargazers

Watchers

alphagan-matting's Issues

另外一个代码

我看到 medium上有篇文章解释了这个论文，然后甩出了另外一个代码 pytorch alpha gan matting

background images data set size

@CDOTAD Hi, I have applied for the Adobe data set, Is your data set processing combined with COCO data ?

Could you offer pretrained model?

关于训练的讨论

你好，最近我也在这个模型去做课题，训练也碰上一些GD的平衡问题，可以留个联系方式讨论下吗？：）

Test: Clip vs. Non-clip

Hi, Good job!
I have been working on this paper as well. I found the same issue during test where there are obvious edge discontinuities at the clip borders! For the case where you did not use clip, did you just down-sample the original image to 320*320, and resize the alpha matte back to the original size? I found down-sample works pretty bad on high-resolution images :( Did you find a better way to run test on high-resolution images by any chance? Thank you.

WARNING:root:Setting up a new session...
INFO:tornado.access:200 POST /env/alphaGAN (127.0.0.1) 0.28ms
WARNING:visdom:Without the incoming socket you cannot receive events from the server or register event handlers to your Visdom client.
0it [00:00, ?it/s]Traceback (most recent call last):
  File "alphaGAN_train.py", line 80, in <module>
    main()
  File "alphaGAN_train.py", line 76, in main
    gan.train(dataset)
  File "/home/xxx/work/AlphaGANMatting/model/AlphaGAN.py", line 462, in train
    for ii, data in tqdm.tqdm(enumerate(dataset)):
  File "/usr/local/anaconda3/lib/python3.6/site-packages/tqdm/_tqdm.py", line 1022, in __iter__
    for obj in iterable:
  File "/home/xxx/work/AlphaGANMatting/data/__init__.py", line 23, in __iter__
    for i, data in enumerate(self.dataloader):
  File "/usr/local/anaconda3/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 637, in __next__
    return self._process_next_batch(batch)
  File "/usr/local/anaconda3/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 658, in _process_next_batch
    raise batch.exc_type(batch.exc_msg)
RuntimeError: Traceback (most recent call last):
  File "/usr/local/anaconda3/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 138, in _worker_loop
    samples = collate_fn([dataset[i] for i in batch_indices])
  File "/usr/local/anaconda3/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 138, in <listcomp>
    samples = collate_fn([dataset[i] for i in batch_indices])
  File "/home/xxx/work/AlphaGANMatting/data/input_dataset.py", line 98, in __getitem__
    T = self.transform(trimap_img)
  File "/usr/local/anaconda3/lib/python3.6/site-packages/torchvision/transforms/transforms.py", line 60, in __call__
    img = t(img)
  File "/usr/local/anaconda3/lib/python3.6/site-packages/torchvision/transforms/transforms.py", line 163, in __call__
    return F.normalize(tensor, self.mean, self.std, self.inplace)
  File "/usr/local/anaconda3/lib/python3.6/site-packages/torchvision/transforms/functional.py", line 208, in normalize
    tensor.sub_(mean[:, None, None]).div_(std[:, None, None])
RuntimeError: output with shape [1, 2278, 3138] doesn't match the broadcast shape [3, 2278, 3138]

我从http://alphamatting.com/datasets.php 这里下载的原图，请问是什么地方需要修改吗？

报错的相关代码

  input_img = Image.open(input_path).convert('RGB')
        trimap_img = Image.open(trimap_path)
        alpha_img = Image.open(alpha_path)
        bg_img = Image.open(bg_path).convert('RGB')
        fg_img = Image.open(fg_path).convert('RGB')

        #x, y = random_choice(trimap_img)

        I = self.transform(input_img)
        T = self.transform(trimap_img)
        A = self.transform(alpha_img)
        B = self.transform(bg_img)
        F = self.transform(fg_img)