Skip to content

Latest commit

 

History

History
73 lines (72 loc) · 4.19 KB

training_with_your_own_dataset.md

File metadata and controls

73 lines (72 loc) · 4.19 KB

Training with your own dataset

Training with your own dataset requires a GPU with 6G memory (above GTX1060).
We will make "face" as an example. If you don't have any picture, you can download CelebA or WIDER.

Getting Started

Prerequisites

Dependencies

This code depends on opencv-python, torchvision, matplotlib available via pip install.

Clone this repo

git clone https://github.com/HypoX64/DeepMosaics
cd DeepMosaics

Make training datasets

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.

  • You can use draw_mask.pyto generate them.
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
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)
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)
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)
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)
python make_video_dataset.py --datadir 'dir for your videos' --model_path ../pretrained_models/mosaic/add_face.pth --mask_threshold 96 --savedir ../datasets/video/face

Training

Add

cd train/add
python train.py --use_gpu 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 or pix2pixHD to train model. We just take pix2pixHD as an example.
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)
cd train/clean
python train.py --dataset ../../datasets/video/face --savename face --savefreq 100000 --gan --hd --lr 0.0002 --lambda_gan 1 --use_gpu 0

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