import gradio as gr
import requests
import base64
import os

def image_to_base64(image):
    with open(image.name, "rb") as file:
        encoded_string = base64.b64encode(file.read())
        return encoded_string.decode("utf-8")

def anti_nsfw(image):
    allowed_extensions = [".png", ".jpeg", ".jpg"]
    file_extension = os.path.splitext(image.name)[1].lower()
    
    if file_extension not in allowed_extensions:
        return "Invalid file type. Please upload a PNG or JPEG image."
    
    encoded_image = image_to_base64(image)
    resp = requests.post(
        "https://8d2e08d2-7720-461c-8281-4ad7f69557e7.id.repl.co/scan",
        data={"base64": f"data:image/png;base64,{encoded_image}"}
    )
    return resp.text

iface = gr.Interface(
    fn=anti_nsfw,
    inputs="file",
    outputs="text",
    title="AntiNSFW",
    description="Check if an image is safe for work (SFW) or not safe for work (NSFW)."
)

iface.launch()