nsfw_detection.py

from PIL import Image
import requests
from io import BytesIO
import base64
import re
from transformers import pipeline

classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")


def classify_image_if_nsfw(image_url: str):
    try:
        # Check if it's a base64 data URL
        if image_url.startswith('data:image'):
            print("Processing base64 data URL")

            # Extract the base64 data from the data URL
            match = re.match(r'data:image/(?P<ext>\w+);base64,(?P<data>.*)', image_url)
            if not match:
                raise ValueError("Invalid base64 data URL format")

            base64_data = match.group('data')
            image_format = match.group('ext')

            # Decode the base64 data
            image_data = base64.b64decode(base64_data)

            # Open the image from decoded data
            img = Image.open(BytesIO(image_data))

        else:
            # It's a regular URL - download the image
            print("Processing regular URL")
            response = requests.get(image_url)
            response.raise_for_status()

            # Open and process the image
            img = Image.open(BytesIO(response.content))

        print("Image size:", img.size)
        print("Image format:", img.format)
        print("Image mode:", img.mode)

        # Ensure image is in RGB mode (required by most models)
        if img.mode != 'RGB':
            img = img.convert('RGB')

        # Classify the image
        classifier_response = classifier(img)
        print("Classifier Response:", classifier_response)

        return classifier_response

    except Exception as e:
        print(f"Error processing image: {e}")
        raise

# Example usage with both types:
# Regular URL
# result1 = classify_image_if_nsfw("https://example.com/image.jpg")

# Base64 data URL (you would use an actual base64 string here)
# result2 = classify_image_if_nsfw("data:image/jpeg;base64,/9j/4AAQSkZJRgABAQ...")