NSFW classify model implemented with tensorflow. Use nsfw dataset provided here https://github.com/alexkimxyz/nsfw_data_scraper Thanks for sharing the dataset with us. You can find all the model details here. Don not hesitate to raise an issue if you're confused with the model.
This software has only been tested on ubuntu 16.04(x64). Here is the test environment info
OS: Ubuntu 16.04 LTS
GPU: Two GTX 1070TI
CUDA: cuda 9.0
Tensorflow: tensorflow 1.12.0
OPENCV: opencv 3.4.1
NUMPY: numpy 1.15.1
Other required package you may install them by
pip3 install -r requirements.txt
In this repo I uploaded a pretrained dataset introduced as before, you need to download the pretrained weights_file in folder REPO_ROOT_DIR/model/.
You can test a single image on the trained model as follows
cd REPO_ROOT_DIR
python tools/test_model.py --weights_path model/new_model/nsfw_cls.ckpt-100000
--image_path data/test_data/test_drawing.jpg
The following part will show you how the dataset is well prepared
First you need to download all the origin nsfw data. Here is the how_to_download_source_data. The training example should be organized like the what you can see in REPO_ROOT_DIR/data/nsfw_dataset_example. Then you should modified the REPO_ROOT_DIR/tools/make_nsfw_dataset.sh with your local nsfw dataset. Then excute the dataset preparation script. That may take about one hour in my local machine. You may enlarge the __C.TRAIN.CPU_MULTI_PROCESS_NUMS in config/global_config.py if you have a powerful cpu to accelerate the prepare process.
cd REPO_ROOT_DIR
bash tools/make_nsfw_dataset.sh
The image of each subclass will be split into three part according to the ratio training : validation : test = 0.75 : 0.1 : 0.15. All the image will be convert into tensorflow format record for efficient importing data pipline.
The model support multi-gpu training. If you want to training the model on multiple gpus you need to first adjust the __C.TRAIN.GPU_NUM in config/global_config.py file. Then excute the multi-gpu training procedure as follows:
cd REPO_ROOT_DIR
python tools/train_nsfw.py --dataset_dir PATH/TO/PREPARED_NSFW_DATASET --use_multi_gpu True
If you want to train the model from last snap shot you may excute following command:
cd REPO_ROOT_DIR
python tools/train_nsfw.py --dataset_dir PATH/TO/PREPARED_NSFW_DATASET
--use_multi_gpu True --weights_path PATH/TO/YOUR/LAST_CKPT_FILE_PATH
You may set the --use_multi_gpu False then the whole training process will be excuted on single gpu.
The main model's hyperparameter are as follows:
iterations nums: 160010
learning rate: 0.1
batch size: 32
origin image size: 256
cropped image size: 224
training example nums: 159477
testing example nums: 31895
validation example nums: 21266
The rest of the hyperparameter can be found here.
If you want to convert the downloaded ckpt model into tensorflow saved model you can simply modify the file path in ROOT_DIR/tools/export_nsfw_saved_model.sh and run it.
bash tools/export_nsfw_saved_model.sh
You may monitor the training process using tensorboard tools
During my experiment the train loss drops as follows:
The train_top_1_error drops as follows:
The validation loss drops as follows:
The validation_top_1_error drops as follows:
You can evaluate the model's performance on the nsfw dataset prepared in advance as follows
cd REPO_ROOT_DIR
python tools/evaluate_nsfw.py --weights_path model/new_model/nsfw_cls.ckpt-160000
--dataset_dir PATH/TO/YOUR/NSFW_DATASET
After you run the script you should see something like this
The model's main evaluation index are as follows:
Precision: 0.92406 with average weighted on each class
Recall: 0.92364 with average weighted on each class
F1 score: 0.92344 with average weighted on each class
The Confusion_Matrix is as follows:
The Precison_Recall is as follows:
Since tensorflo-js is well supported the online deep learning is easy to deploy. Here I have make a online demo to do local nsfw classification work. The whole js work can be found here https://github.com/MaybeShewill-CV/MaybeShewill-CV.github.io/tree/master/nsfw_classification I have supplied a tool to convert the trained tensorflow saved model file into tensorflow js model file. In order to generate saved model you can read the description about it above. After you generate the tensorflow saved model you can simply modify the file path and run the following script
cd ROOT_DIR
bash tools/convert_tfjs_model.sh
The online demo's example are as follows:
- Add tensorflow serving script