menu/get_menu/generate_menu_icons.py at main · engie/menu · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
"""
generate_menu_icons.py

Generates (or reuses) 100×100 black‑and‑white vector‑style icons for
each menu item in a school term menu JSON file.  Adds a base64‑encoded
version of each icon back into the JSON (as meat_logo, meatfree_logo,
dessert_logo) and writes the augmented JSON out.

Usage:
    export OPENAI_API_KEY="sk‑..."
    python generate_menu_icons.py term_menu.json term_menu_with_icons.json

Dependencies:
    pip install openai pillow python‑slugify

© 2025
"""

import json
import base64
import re
import os
from dotenv import load_dotenv
load_dotenv()
import sys
from pathlib import Path
from io import BytesIO
import argparse

import openai
from PIL import Image
from slugify import slugify

# ---------- Configuration ----------
ICON_DIR = Path("icons")          # Where icons live on disk
ICON_SIZE = (100, 100)            # Pixels (width, height)
OPENAI_IMAGE_MODEL = "dall-e-3"   # Or another available model
# -----------------------------------

PROMPT_TEMPLATE = (
    "Black-and-white vector icon: {food}. "
    "Bold, clean outlines, no shading or gray, designed for a 100x100 pixel square canvas."
)

logo_book = {}

def ensure_icon_in_book(food_string: str) -> str:
    """Ensure an icon exists on disk at full resolution and
    return a *100×100* resized copy in memory (PIL Image)."""
    ICON_DIR.mkdir(exist_ok=True)
    slug = slugify(food_string, lowercase=True, separator="_")
    icon_path = ICON_DIR / f"{slug}.png"

    if slug in logo_book:
        return slug

    if not icon_path.exists():
        print(f"Generating icon for '{food_string}' → {icon_path}")
        # ---- Generate with OpenAI images API ----
        response = openai.images.generate(
            model=OPENAI_IMAGE_MODEL,
            prompt=PROMPT_TEMPLATE.format(food=food_string),
            n=1,
            size="1024x1024",  # Generate larger then downscale
            response_format="b64_json",
        )
        b64_png = response.data[0].b64_json
        raw_png = base64.b64decode(b64_png)
        Image.open(BytesIO(raw_png)).save(icon_path)

    # Always leave the on‑disk file at its original size; create a resized copy in memory
    with Image.open(icon_path) as im:
        im = im.convert("RGBA").convert('1', dither=Image.NONE)
        resized = im.resize(ICON_SIZE, resample=Image.NEAREST)
        logo_book[slug] = icon_as_base64(resized)
    return slug


def icon_as_base64(img: Image.Image) -> str:
    """Return base64 PNG string (no data: prefix) for the given in‑memory PIL Image."""
    buf = BytesIO()
    img.save(buf, format="PNG", optimize=True)
    return base64.b64encode(buf.getvalue()).decode("ascii")


def process_menu(in_path: Path, out_path: Path, debug: bool = False):
    with open(in_path, "r", encoding="utf-8") as f:
        menu = json.load(f)

    for date, items in menu.items():
        for key in ("meat", "meatfree", "desert"):
            if key in items:
                logo_key = f"{key}_logo"
                if logo_key not in items:
                    items[logo_key] = ensure_icon_in_book(items[key])

    # Write augmented JSON
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump({
            "dates": menu,
            "logos": logo_book,
        }, f, ensure_ascii=False, indent=2)

    if debug:
        html_path = out_path.with_suffix(".html")
        with open(html_path, "w", encoding="utf-8") as f:
            f.write("<html><head><title>Icon Preview</title></head><body>\n")
            f.write("<h1>Generated Icons</h1>\n")
            for slug, b64 in logo_book.items():
                f.write(f"<div style='display:inline-block; text-align:center; margin:10px'><img src='data:image/png;base64,{b64}' width='100' height='100'><br>{slug}</div>\n")
            f.write("</body></html>\n")
        print(f"Debug HTML written to → {html_path}")

    print(f"Wrote updated menu with icons → {out_path}")


def main(argv=None):
    argv = argv or sys.argv[1:]
    parser = argparse.ArgumentParser(description="Generate menu icons and embed base64 data.")
    parser.add_argument("in_json", help="Input JSON file")
    parser.add_argument("out_json", help="Output JSON file")
    parser.add_argument("--debug", action="store_true", help="Emit HTML preview of icons")
    args = parser.parse_args(argv)

    in_path = Path(args.in_json)
    out_path = Path(args.out_json)

    process_menu(in_path, out_path, debug=args.debug)


if __name__ == "__main__":
    main()