szuru-eink-bot/szuru-eink.py

#!/usr/bin/python3

"""Bot that uploads a random szurubooru image to an eink screen"""

import random
import requests
import os
import pyszuru
from dotenv import load_dotenv
from PIL import Image

load_dotenv()

from epd.epd_config import palArr, epdArr, EPD_TYPES
from epd.image_processor import getErr, getNear, addVal, procImg
from epd.epd_uploader import upload_image


# get truly random post
try:
    szuru_url = os.environ["SZURU_URL"]
    szuru_user = os.environ["SZURU_USER"]
    szuru_token = os.environ["SZURU_TOKEN"]
except KeyError as e:
    raise ValueError(f"Missing required environment variable: {e}")

booru = pyszuru.API(szuru_url, username=szuru_user, token=szuru_token)
highest_post = next(booru.search_post("sort:id type:image", page_size=1))

post = None
while post is None:
    random_id = random.randint(0, highest_post.id_)
    temp_post = booru.getPost(random_id)
    if temp_post and temp_post.mime.startswith("image/") and temp_post.content:
        post = temp_post
        print(f"[DEBUG] Found image post with ID: {post.id_}, MIME: {post.mime}")
    else:
        print(
            f"[DEBUG] Skipping post ID: {random_id} (MIME: {temp_post.mime if temp_post else 'None'}, Content: {bool(temp_post.content) if temp_post else 'None'}) - not a valid image."
        )

# download image
image_downloaded = False
while not image_downloaded:
    try:
        r = requests.get(post.content, stream=True)
        print(f"[DEBUG] Downloading image from {post.content}")
        r.raise_for_status()
        with open("image.jpg", "wb") as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)
        print("[DEBUG] Image downloaded and saved as image.jpg")

        # Attempt to open the image to verify
        img = Image.open("image.jpg")
        img.verify()  # Verify that it is an image
        img.close()
        image_downloaded = True
        print("[DEBUG] Image successfully verified.")
    except (PIL.UnidentifiedImageError, requests.exceptions.RequestException) as e:
        print(
            f"[ERROR] Failed to process downloaded file (Error: {e}). Retrying with a new post..."
        )
        # Clean up the invalid file
        if os.path.exists("image.jpg"):
            os.remove("image.jpg")
        # Find a new post
        post = None
        while post is None:
            random_id = random.randint(0, highest_post.id_)
            temp_post = booru.getPost(random_id)
            if temp_post and temp_post.mime.startswith("image/") and temp_post.content:
                post = temp_post
                print(
                    f"[DEBUG] Found new image post with ID: {post.id_}, MIME: {post.mime}"
                )
            else:
                print(
                    f"[DEBUG] Skipping post ID: {random_id} (MIME: {temp_post.mime if temp_post else 'None'}, Content: {bool(temp_post.content) if temp_post else 'None'}) - not a valid image."
                )

    # Process and upload
    img = Image.open("image.jpg")

# Get configuration from environment variables
try:
    epd_type_str = os.environ["EPD_TYPE"]
    ip_addr = os.environ["EINK_IP"]
    dithering_mode = os.environ["DITHERING_MODE"].lower()
except KeyError as e:
    raise ValueError(f"Missing required environment variable: {e}")

if epd_type_str not in EPD_TYPES:
    raise ValueError(
        f"Invalid EPD_TYPE: {epd_type_str}. Must be one of {list(EPD_TYPES.keys())}"
    )
epd_ind = EPD_TYPES[epd_type_str]

is_red = False
if dithering_mode == "color":
    # Check if the selected EPD type supports color based on scriptD.js logic
    # EPDs that are explicitly monochrome in scriptD.js despite palArr:
    monochrome_only_epds = [0, 3, 6, 7, 9, 12, 16, 19, 22, 26, 27, 28, 39, 40, 43]
    if epd_ind in monochrome_only_epds:
        print(
            f"[WARNING] Dithering mode is set to 'color', but EPD type '{epd_type_str}' (Index {epd_ind}) is treated as monochrome by the device's protocol. Color will not appear."
        )
        is_red = False  # Force monochrome for this EPD type
    else:
        is_red = True
elif dithering_mode != "mono":
    raise ValueError(
        f"Invalid DITHERING_MODE: {dithering_mode}. Must be 'mono' or 'color'."
    )

# Get display dimensions and palette from epdArr
dW, dH, pal_ind = epdArr[epd_ind]
target_size = (dW, dH)

# Resize image to fit display, preserving aspect ratio
img.thumbnail(target_size, Image.Resampling.LANCZOS)

# Create a new image with a white background
background = Image.new("RGB", target_size, (255, 255, 255))

# Calculate position and paste the thumbnail onto the background
paste_x = (target_size[0] - img.width) // 2
paste_y = (target_size[1] - img.height) // 2
background.paste(img, (paste_x, paste_y))
img = background

print(f"[DEBUG] Dithering Mode: {dithering_mode}")
print(f"[DEBUG] is_red: {is_red}")
print(f"[DEBUG] EPD Type String: {epd_type_str}")
print(f"[DEBUG] EPD Index: {epd_ind}")
print(f"[DEBUG] Palette Index: {pal_ind}")
# Process the final image with dithering
processed_img = procImg(
    img, dW, dH, 0, 0, pal_ind, isLvl=False, isRed=is_red, palArr=palArr
)

processed_img.save("processed.png")  # Save for debugging
print("[DEBUG] Imaged processed and saved as processed.png")

upload_image(processed_img, ip_addr, epd_ind, epdArr, palArr, getNear)