Introduction
Detecting Not-Suitable-For-Work (NSFW) images is a high demand task in computer vision. While there are many types of NSFW images, here we focus on the pornographic images.
The Yahoo Open-NSFW model originally developed with the Caffe framework has been a favourite choice, but the work is now discontinued and Caffe is also becoming less popular. Please see the description on the Yahoo project page for the context, definitions, and model training details.
This Open-NSFW 2 project provides a TensorFlow 2 implementation of the Yahoo model, with references to its previous third-party TensorFlow 1 implementation.
Installation
Python 3.7 or above is required. Tested for 3.7, 3.8, and 3.9.
The best way to install Open-NSFW 2 with its dependencies is from PyPI:
python3 -m pip install --upgrade opennsfw2
Alternatively, to obtain the latest version from this repository:
git clone [email protected]:bhky/opennsfw2.git
cd opennsfw2
python3 -m pip install .
Usage
import numpy as np
import opennsfw2 as n2
from PIL import Image
# Load and preprocess image.
image_path = "path/to/your/image.jpg"
pil_image = Image.open(image_path)
image = n2.preprocess_image(pil_image, n2.Preprocessing.YAHOO)
# The preprocessed image is a NumPy array of shape (224, 224, 3).
# Create the model.
# By default, this call will search for the pre-trained weights file from path:
# $HOME/.opennsfw2/weights/open_nsfw_weights.h5
# If not exists, the file will be downloaded from this repository.
# The model is a `tf.keras.Model` object.
model = n2.make_open_nsfw_model()
# Make predictions.
inputs = np.expand_dims(image, axis=0) # Add batch axis (for single image).
predictions = model.predict(inputs)
# The shape of predictions is (batch_size, 2).
# Each row gives [sfw_probability, nsfw_probability] of an input image, e.g.:
sfw_probability, nsfw_probability = predictions[0]
Alternatively, the end-to-end pipeline function can be used:
import opennsfw2 as n2
image_paths = [
"path/to/your/image1.jpg",
"path/to/your/image2.jpg",
# ...
]
predictions = n2.predict(
image_paths, batch_size=4, preprocessing=n2.Preprocessing.YAHOO
)
API
preprocess_image
Apply necessary preprocessing to the input image.
- Parameters:
pil_image
(PIL.Image
): Input as a Pillow image.preprocessing
(Preprocessing
enum, defaultPreprocessing.YAHOO
): See preprocessing details.
- Return:
- NumPy array of shape
(224, 224, 3)
.
- NumPy array of shape
Preprocessing
Enum class for preprocessing options.
Preprocessing.YAHOO
Preprocessing.SIMPLE
make_open_nsfw_model
Create an instance of the NSFW model, optionally with pre-trained weights from Yahoo.
- Parameters:
input_shape
(Tuple[int, int, int]
, default(224, 224, 3)
): Input shape of the model, this should not be changed.weights_path
(Optional[str]
, default$HOME/.opennsfw/weights/open_nsfw_weights.h5
): Path to the weights in HDF5 format to be loaded by the model. The weights file will be downloaded if not exists. Users can provide path if the default is not preferred. IfNone
, no weights will be downloaded nor loaded to the model.
- Return:
tf.keras.Model
object.
predict
End-to-end pipeline function from input image paths to predictions.
- Parameters:
image_paths
(Sequence[str]
): List of paths to input image files.batch_size
(int
, default32
): Batch size to be used for model inference.preprocessing
: Same as that inpreprocess_image
.weights_path
: Same as that inmake_open_nsfw_model
.
- Return:
- NumPy array of shape
(batch_size, 2)
, each row gives[sfw_probability, nsfw_probability]
of an input image.
- NumPy array of shape
Preprocessing details
Options
This implementation provides the following preprocessing options.
YAHOO
: The default option which was used in the original Yahoo's Caffe and the later TensorFlow 1 implementations. The key steps are:- Resize the input Pillow image to
(256, 256)
. - Save the image as JPEG bytes and reload again to a NumPy image (this step is mysterious, but somehow it really makes a difference).
- Crop the centre part of the NumPy image with size
(224, 224)
. - Swap the colour channels to BGR.
- Subtract the training dataset mean value of each channel:
[104, 117, 123]
.
- Resize the input Pillow image to
SIMPLE
: A simpler and probably more intuitive preprocessing option is also provided, but note that the model output probabilities will be different. The key steps are:- Resize the input Pillow image to
(224, 224)
. - Convert to a NumPy image.
- Swap the colour channels to BGR.
- Subtract the training dataset mean value of each channel:
[104, 117, 123]
.
- Resize the input Pillow image to
Comparison
Using 521 private images, the NSFW probabilities given by three different settings are compared:
- TensorFlow 1 implementation with
YAHOO
preprocessing. - TensorFlow 2 implementation with
YAHOO
preprocessing. - TensorFlow 2 implementation with
SIMPLE
preprocessing.
The following figure shows the result:
The current TensorFlow 2 implementation with YAHOO
preprocessing can totally reproduce the well-tested TensorFlow 1 result, with small floating point errors only.
With SIMPLE
preprocessing the results are different, where the model tends to give lower probabilities.