app.py

import os
from json import dumps
import gradio as gr
from huggingface_hub import login
from transformers import pipeline

from classify_utils import classify_image_if_nsfw, get_nsfw_classifier, get_normal_classifier


def init_login():
    try:
        login(os.environ.get('HF_TOKEN'))
        print("logged in successfully")
    except Exception as e:
        print("issue logging into hugging face:", e)


def init_classifier():
    global classifier
    classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")



def respond(
    message,
    history: list[dict[str, str]],
):
    json_response = classify_image_if_nsfw(classifier, message)
    print(dumps(json_response, indent=4))
    normal_classifier = get_normal_classifier(json_response)
    nsfw_classifier = get_nsfw_classifier(json_response)
    # Fixed: Use 'and' instead of '&&', and improved logic flow
    if nsfw_classifier is not None and nsfw_classifier['score'] > 0.75:
        return {"text": "Very Explicit"}
    elif nsfw_classifier is not None and nsfw_classifier['score'] > 0.5:
        return {"text": "Somewhat Explicit"}
    elif normal_classifier is not None and normal_classifier['score'] < 0.8:
        return {"text": "Somewhat Normal"}
    else:
        return {"text": "Normal"}



"""
Functionality of nsfw detector
"""
init_login()
init_classifier()

chatbot = gr.ChatInterface(
    respond,
    type="messages",
    additional_inputs=[]
)

with gr.Blocks() as demo:
    chatbot.render()

if __name__ == "__main__":
    demo.launch()