speckles/main.py

import io
from typing import Literal
import itertools
import logging
import random

import matplotlib.pyplot as plt
import numpy as np
import uvicorn
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse

app = FastAPI(title="Speckles API", root_path="/images")
origins = [
    "http://localhost",
    "http://localhost:3000",
    "https://localhost",
    "https://0124816.xyz",
    "http://0124816.xyz:3001",
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

MEDIA_TYPES = {
    "png": "image/png",
    "jpeg": "image/jpeg",
    "jpg": "image/jpeg",
    "svg": "image/svg+xml",
    "pdf": "application/pdf",
}


Marker = Literal[
    ".",
    ",",
    "o",
    "v",
    "^",
    "<",
    ">",
    "1",
    "2",
    "3",
    "4",
    "8",
    "s",
    "p",
    "P",
    "*",
    "h",
    "H",
    "+",
    "x",
    "X",
    "D",
    "d",
    "|",
    "_",
]


@app.get("/speckles/")
def make_wallpaper(
    speckle_colours: str,
    density: float | None = 0.12,
    size: float | None = 3,
    fileformat: str = "svg",
    orientation: str | None = "landscape",
    local: bool = False,
    marker: str | None = ".",
):
    if fileformat not in MEDIA_TYPES:
        return
    speckle_colours = speckle_colours.split(",")
    background = speckle_colours.pop(0)
    if orientation == "portrait":
        x, y = (1080, 1920)
    elif orientation == "landscape":
        x, y = (1920, 1080)
    elif "x" in orientation:
        resolution = orientation.split("x")
        if len(resolution) != 2:
            logging.critical("input resolution has more or less than 2 dimensions")
            return
        x, y = resolution
        if all([x.isdigit(), y.isdigit()]):
            x, y = int(x), int(y)
        else:
            return
    else:
        x, y = (1920, 1080)
    speckles_per_colour = int(x / 128 * y / 128 * density)

    fig, ax = plt.subplots(figsize=(x / 120, y / 120), facecolor=background)
    ax.set_facecolor(background)
    [spine.set_color(background) for spine in ax.spines.values()]
    ax.set_xticks([])
    ax.set_yticks([])
    ax.margins(0, 0)

    for color, size in itertools.product(
        speckle_colours,
        np.logspace(0, size, 10, base=np.exp(2)),
    ):
        ax.scatter(
            [random.random() * x / 8 for _ in range(speckles_per_colour)],
            [random.random() * y / 8 for _ in range(speckles_per_colour)],
            c=color,
            s=size,
            marker=random.choice(marker),
        )

    fig.tight_layout()
    # plt.xlim(0, x)
    # plt.ylim(0, y)
    # plt.axis("off")

    buf = io.BytesIO()
    fig.savefig(
        buf,
        format=fileformat,
        dpi=128,
        bbox_inches="tight",
        pad_inches=0,
    )
    buf.seek(0)
    if local:
        filename = f"speckles.{fileformat}"
        with open(filename, "wb") as f:
            f.write(buf.getbuffer())
        return filename
    else:
        return StreamingResponse(content=buf, media_type=MEDIA_TYPES[fileformat])
    buf.close()


if __name__ == "__main__":
    uvicorn.run("main:app", workers=2, port=8099, reload=False)