189 lines
6.8 KiB
Python
189 lines
6.8 KiB
Python
"""
|
|
|
|
AUTHOR:
|
|
|
|
Khushal P Soonderji
|
|
|
|
DATE:
|
|
|
|
Monday, 2nd Sept., 2024
|
|
|
|
OBJECTIVE:
|
|
|
|
To provide an easy way to assess images for adult content.
|
|
Tried and implemented using HuggingFace models.
|
|
|
|
REFERENCES:
|
|
|
|
N/A
|
|
|
|
DOWNLOADS:
|
|
|
|
N/A
|
|
|
|
"""
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** IMPORT ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# To make sibling directories accessible for imports:
|
|
import sys
|
|
sys.path.append(".")
|
|
sys.path.append("..")
|
|
|
|
# System-level activities:
|
|
import io
|
|
|
|
# For using the AI model:
|
|
import torch
|
|
from transformers import pipeline
|
|
|
|
# To download images from URLs:
|
|
import requests
|
|
|
|
# To read images:
|
|
from PIL import Image
|
|
|
|
# Common:
|
|
from shared import variables
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** MACROS / ONE-TIME INIT ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** VARIABLES ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** CLASSES ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
class AssessImageNSFW:
|
|
|
|
def __init__(
|
|
self,
|
|
model,
|
|
nsfw_label = "nsfw",
|
|
nsfw_threshold = 0.25,
|
|
):
|
|
|
|
"""
|
|
This class assesses the input image and tells if it is blurry or clear.
|
|
:param model: The model to use. Either the name of the HuggingFace repo, or the directory where the model is
|
|
stored.
|
|
:param nsfw_threshold: The max allowed blurriness (0 to 1 range).
|
|
"""
|
|
|
|
# Note down the config:
|
|
self.__model = model
|
|
self.__device = "gpu" if torch.cuda.is_available() else "cpu"
|
|
self.__nsfw_label = nsfw_label
|
|
self.__nsfw_threshold = nsfw_threshold
|
|
|
|
# Initialize the classifier:
|
|
self.__classifier = pipeline(
|
|
task = "image-classification",
|
|
model = self.__model,
|
|
device = self.__device
|
|
)
|
|
|
|
@staticmethod
|
|
async def read_image(source):
|
|
|
|
"""
|
|
Reads an image in whichever format it is provided and returns it as a PIL object.
|
|
:param source: The image as either a path or a URL or a io.BytesIO object.
|
|
:return: The image opened as a PIL object.
|
|
"""
|
|
|
|
# If the input image is already a PIL image:
|
|
if isinstance(source, Image.Image): return source
|
|
|
|
# If a buffer is provided:
|
|
elif isinstance(source, io.BytesIO):
|
|
source.seek(0)
|
|
return Image.open(source)
|
|
|
|
# If a string is provided (local path or URL):
|
|
elif isinstance(source, str):
|
|
if source.startswith("http://") or source.startswith("https://"):
|
|
response = await variables.http_client.get(source)
|
|
return Image.open(io.BytesIO(response.content))
|
|
else: return Image.open(source)
|
|
|
|
async def classify(self, image):
|
|
|
|
"""
|
|
To get the prediction of the model from the given input image.
|
|
:param image: The image (as a PIL object or file path or a URL).
|
|
:return: The dictionary of classes with their respective predictions.
|
|
"""
|
|
|
|
image = await self.read_image(image)
|
|
predictions = self.__classifier(image)
|
|
classes = {p["label"]: p["score"] for p in predictions}
|
|
return classes
|
|
|
|
async def is_ok(
|
|
self,
|
|
image,
|
|
nsfw_threshold = None
|
|
):
|
|
|
|
"""
|
|
Checks if the image given to it can be used, or should be rejected.
|
|
:param image: The image (as a PIL object or file path or a URL).
|
|
:param nsfw_threshold: A custom threshold to test against. If not provided, the default will be taken that was
|
|
provided when the instance was created.
|
|
:return: True if the image is okay, else False.
|
|
"""
|
|
|
|
image = await self.read_image(image)
|
|
nsfw_threshold = nsfw_threshold or self.__nsfw_threshold
|
|
classes = await self.classify(image)
|
|
return True if classes[self.__nsfw_label] <= nsfw_threshold else False
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** MAIN PROGRAM ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
import asyncio
|
|
|
|
|
|
async def main():
|
|
|
|
assessor = AssessImageNSFW(model = r"/path/to/model")
|
|
usable = await assessor.is_ok(r"/path/to/image/img.jpg")
|
|
classes = await assessor.classify(r"https://...")
|
|
|
|
|
|
asyncio.run(main())
|