# Training with your own dataset
Training with your own dataset requires a GPU with 6G memory (above GTX1060).<br>
We will make "face" as an example. If you don't have any picture, you can download [CelebA](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html) or [WIDER](http://mmlab.ie.cuhk.edu.hk/projects/WIDERFace/WiderFace_Results.html).

## Getting Started
#### Prerequisites
  - Linux, Mac OS, Windows
  - Python 3.6+
  - [ffmpeg 3.4.6](http://ffmpeg.org/)
  - [Pytorch 1.0+](https://pytorch.org/)
  - NVIDIA GPU(with more than 6G memory) + CUDA CuDNN<br>
#### Dependencies
This code depends on opencv-python, torchvision, matplotlib, tensorboardX, scikit-image available via conda install.
```bash
# or
pip install -r requirements.txt
```
#### Clone this repo
```bash
git clone https://github.com/HypoX64/DeepMosaics
cd DeepMosaics
```
## Make training datasets
```bash
cd make_datasets
```
### Add mosaic dataset
Please generate mask from images which you want to add mosaic(number of images should be above 1000). And then put the images in ```face/origin_image```, and masks in ```face/mask```.<br>
* You can use ```draw_mask.py```to generate them.
```bash
python draw_mask.py --datadir 'dir for your pictures' --savedir ../datasets/draw/face
#Press the left mouse button to draw the mask .  Press 'S' to save mask, 'A' to reduce  brush size, 'D' to increase brush size, 'W' to cancel drawing.
```
* If you want to get images from videos, you can use ```get_image_from_video.py```
```bash
python get_image_from_video.py --datadir 'dir for your videos' --savedir ../datasets/video2image --fps 1
```
### Clean mosaic dataset
We provide several methods for generating clean mosaic datasets. However, for better effect, we recommend train a addmosaic model in a small data  first and use it to automatically generate datasets in a big data. (recommend: Method 2(for image) & Method 4(for video))
* Method 1: Use drawn mask to make pix2pix(HD) datasets   (Require``` origin_image``` and ```mask```)
```bash
python make_pix2pix_dataset.py --datadir ../datasets/draw/face --hd --outsize 512 --fold 1 --name face --savedir ../datasets/pix2pix/face --mod drawn --minsize 128 --square
```
* Method 2: Use addmosaic model to make pix2pix(HD) datasets (Require addmosaic pre-trained model)
```bash
python make_pix2pix_dataset.py --datadir 'dir for your pictures' --hd --outsize 512 --fold 1 --name face --savedir ../datasets/pix2pix/face --mod network --model_path ../pretrained_models/mosaic/add_face.pth --minsize 128 --square --mask_threshold 128
```
* Method 3: Use Irregular Masks to make pix2pix(HD) datasets (Require [Irregular Masks](https://nv-adlr.github.io/publication/partialconv-inpainting))
```bash
python make_pix2pix_dataset.py --datadir 'dir for your pictures' --hd --outsize 512 --fold 1 --name face --savedir ../datasets/pix2pix/face --mod irregular --irrholedir ../datasets/Irregular_Holes_mask --square
```
* Method 4: Use addmosaic model to make video datasets (Require addmosaic pre-trained model. This is better for processing video mosaics)
```bash
python make_video_dataset.py --model_path ../pretrained_models/mosaic/add_face.pth --gpu_id 0 --datadir 'dir for your videos' --savedir ../datasets/video/face
```
## Training
### Add
```bash
cd train/add
python train.py --gpu_id 0 --dataset ../../datasets/draw/face --savename face --loadsize 512 --finesize 360 --batchsize 16
```
### Clean
* For image datasets (generated by ```make_pix2pix_dataset.py```)
We use [pix2pix](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix) or [pix2pixHD](https://github.com/NVIDIA/pix2pixHD) to train model. We just take pix2pixHD as an example.
```bash
git clone https://github.com/NVIDIA/pix2pixHD
cd pix2pixHD
pip install dominate
python train.py --name face --resize_or_crop resize_and_crop --loadSize 563 --fineSize 512 --label_nc 0 --no_instance --dataroot ../datasets/pix2pix/face
```
* For video datasets (generated by ```make_video_dataset.py```)
```bash
cd train/clean
python train.py --dataset ../../datasets/video/face --savename face --n_blocks 4 --lambda_GAN 0.01 --loadsize 286 --finesize 256 --batchsize 16 --n_layers_D 2 --num_D 3 --n_epoch 200 --gpu_id 4,5,6,7 --load_thread 16
```
## Testing
Put saved network to ```./pretrained_models/mosaic/``` and rename it as ```add_face.pth``` or ```clean_face_HD.pth``` or ```clean_face_video_HD.pth```and then run ```deepmosaic.py --model_path ./pretrained_models/mosaic/your_model_name```