bhky / opennsfw2

Keras implementation of the Yahoo Open-NSFW model
MIT License
361 stars 47 forks source link
image-classification image-preprocessing jax keras nsfw-classifier tensorflow2

logo

ci License MIT 1.0

Introduction

Detecting Not-Suitable-For-Work (NSFW) content is a high demand task in computer vision. While there are many types of NSFW content, here we focus on the pornographic images and videos.

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 Keras implementation of the Yahoo model, with references to its previous third-party TensorFlow 1 implementation. Note that Keras 3 is compatible with TensorFlow, JAX, and PyTorch. However, currently this model is only guaranteed to work with TensorFlow and JAX.

A simple API is provided for making predictions on images and videos.

Installation

Tested with TensorFlow and JAX, for Python 3.9 to 3.12.

A note on PyTorch:

The OpenNSFW 2 model can in fact be run on PyTorch, but the biggest issue is that the inference output on PyTorch is quite different from that on TensorFlow and JAX. The reason is still unknown. In addition, inference is much slower on PyTorch probably because of the issues discussed here, i.e., PyTorch uses channels_first for its image data format, but this model uses channels_last (as in TensorFLow and JAX), hence Keras has to convert the channel order back and forth at each layer. Therefore, at the moment it is not recommended to use PyTorch for this model.

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 git@github.com:bhky/opennsfw2.git
cd opennsfw2
python3 -m pip install .

Usage

Quick examples for getting started are given below. For more details, please refer to the API section.

Images

import opennsfw2 as n2

# To get the NSFW probability of a single image, provide your image file path,
# or a `PIL.Image.Image` object.
image_handle = "path/to/your/image.jpg"

nsfw_probability = n2.predict_image(image_handle)

# To get the NSFW probabilities of a list of images, provide a list of file paths,
# or a list of `PIL.Image.Image` objects.
# Using this function is better than looping with `predict_image` as the model 
# will only be instantiated once and batching is done during inference.
image_handles = [
  "path/to/your/image1.jpg",
  "path/to/your/image2.jpg",
  # ...
]

nsfw_probabilities = n2.predict_images(image_handles)

Video

import opennsfw2 as n2

# The video can be in any format supported by OpenCV.
video_path = "path/to/your/video.mp4"

# Return two lists giving the elapsed time in seconds and the NSFW probability of each frame.
elapsed_seconds, nsfw_probabilities = n2.predict_video_frames(video_path)

Lower level with Keras

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 `keras_core.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 (num_images, 2).
# Each row gives [sfw_probability, nsfw_probability] of an input image, e.g.:
sfw_probability, nsfw_probability = predictions[0]

API

preprocess_image

Apply necessary preprocessing to the input image.

Preprocessing

Enum class for preprocessing options.

make_open_nsfw_model

Create an instance of the NSFW model, optionally with pre-trained weights from Yahoo.

predict_image

End-to-end pipeline function from the input image to the predicted NSFW probability.

predict_images

End-to-end pipeline function from the input images to the predicted NSFW probabilities.

Aggregation

Enum class for aggregation options in video frames prediction.

predict_video_frames

End-to-end pipeline function from the input video to predictions.

Preprocessing details

This implementation provides the following preprocessing options.