from openai import OpenAI
import io
import base64
from PIL import Image

client = OpenAI()

def check_img(PIL_image):
    buffer = io.BytesIO()
    PIL_image.save(buffer, format="PNG")
    buffer.seek(0)
    img_bytes = buffer.read()
    base64_bytes = base64.b64encode(img_bytes)
    base64_string = base64_bytes.decode('utf-8')
    base64_string = "data:image/png;base64," + base64_string

    response = client.moderations.create(
        model="omni-moderation-latest",
        input=[
            {
                "type": "image_url",
                "image_url": {
                    "url": base64_string,
                }
            },
        ],
    )

    return response.results[0].flagged

if __name__ == "__main__":
    image = Image.open("lib/11_0M0DY5zO.png")
    flagged = check_img(image)
    print("Flagged:", flagged)