The purpose of the content of this repository is to easily prepare the data to train the Deep Image Matting Network coded by Ge Zheng. The main class (DataCreation2.py) uses the SURREAL Database for the EPS and alpha images, and an Unannotated Database for the backgrounds. However, I have created classes for the main tasks, so you would just need to use PaddingBackgroundCreation.py someone would want to use other sources.
Deep Image Matting paper here
Just run python in the commandline and use the code below. Don't forget to change the folders directory.
from DataCreation2 import DataCreation
# Folder where the surreal data is located
SUR_DATA_DIR = '/data/DataBases/SURREAL/SURREAL/data/cmu/train/run0/01_01'
# Output Directory
OUT_DIR = '/data/HectorSanchez/Deep-Image-Matting/data/'
# Backgrounds
BG_DIR = '/data/HectorSanchez/Deep-Image-Matting/backgrounds/'
create = DataCreation(ipath=SUR_DATA_DIR, opath=OUT_DIR, bg_path=BG_DIR)
create.create_data()There's 4 files that will do the trick.
DataCreation2.py is the main file for generating the data from the SURREAL database.
MatExtractor.py extracts the alpha images from the segm.mat files. Each segm.mat file has a whole lot of alphas
VidMatExtractor.py extracts the correspondig alpha and eps from the segm.mat and .mp4 files.
PaddingBackgroundCreation.py creates the 3 necesary images for training our model (eps, alpha, bg). NOTE: the parameter num_bgs referes to how many backgrounds we are going to use for each eps image. So it randomnly samples num_bgs images from the BG_DIR and combine them with the alpha and eps images. Therefore, if we use num_bgs=100 for each eps/alpha image we are going to have 100 images.
The SURREAL Folder Structure is as follow, but the code above just works at level 4 (run0, run1, etc) and uses the mp4 and segm.mat files.
.
└── data
├── cmu
│ ├── test
│ │ ├── run0
│ │ │ ├── 40_06_c0001_depth.mat
│ │ │ ├── 40_06_c0001_info.mat
│ │ │ ├── 40_06_c0001.mp4
│ │ │ ├── 40_06_c0001_segm.mat
│ │ │ ├── 40_06_c0002_depth.mat
│ │ │ ├── 40_06_c0002_info.mat
│ │ │ ├── 40_06_c0002.mp4
│ │ │ ├── 40_06_c0002_segm.mat
│ │ │ ├── 40_06_c0003_depth.mat
│ │ │ ├── 40_06_c0003_info.mat
│ │ │ ├── 40_06_c0003.mp4
│ │ │ └── 40_06_c0003_segm.mat
│ │ ├── run1
│ │ └── run2
│ └── train
│ ├── run0
│ ├── run1
│ └── run2
└── h36m
├── test
│ ├── run0
│ ├── run1
│ └── run2
└── train
├── run0
├── run1
└── run2
The backgrounds folder just contain all the possible backgrounds to be used in our training
.
├── 983
├── 984
├── 985
├── 986
├── 987
├── 988
├── 989
├── 99
├── 990
├── 991
├── 992
├── 993
├── 994
├── 995
├── 996
├── 997
├── 998
└── 999
The output directory generates the structure shown below:
.
├── alpha1280
│ ├── 0 # folder for alpha image 0
│ │ ├── 0.png # alpha image 0 with background 0
│ │ ├── 1.png # alpha image 0 with background 1
│ │ ├── 2.png # alpha image 0 with background 2
│ │ ├── 3.png # alpha image 0 with background 3
│ │ ├── 4.png # alpha image 0 with background 4
│ │
│ ├── 1 # folder for alpha image 1
│ ├── 2
│ ├── 3
│ ├── 4
│ .
│ .
│ .
│
├── bg
│ ├── 0 # backgrounds for eps/alpha image 0
│ │ ├── 0.png # background 0
│ │ ├── 1.png # background 1
│ │ ├── 2.png # background 2
│ │ ├── 3.png # background 3
│ │ ├── 4.png # background 4
│ │
│ ├── 1
│ ├── 2
│ ├── 3
│ ├── 4
│ .
│ .
│ .
│
└── eps1280
├── 0 # folder for eps image 0
│ ├── 0.png # alpha image 0 with background 0
│ ├── 1.png # alpha image 0 with background 1
│ ├── 2.png # alpha image 0 with background 2
│ ├── 3.png # alpha image 0 with background 3
│ ├── 4.png # alpha image 0 with background 4
│
├── 1
├── 2
├── 3
├── 4
.
.
.