Twarner491
672 lines
29 KiB
Python
Executable File
672 lines
29 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""AvianVisitors - generate kachō-e bird illustrations for a region.
|
|
|
|
Step 1 of the illustration pipeline:
|
|
1. pregen.py render each bird on a uniform cream ground
|
|
2. cutout.py remove the ground (BiRefNet) and crop to the bird
|
|
3. build_masks.py refresh the collage silhouette masks in apt.js
|
|
|
|
Reads a species list (BirdNET-Pi's labels.txt, eBird, or stdin),
|
|
fetches a Wikipedia reference photo for each species, and generates an
|
|
illustration via the Gemini 2.5 Flash Image API. Saves PNGs into
|
|
avian/assets/illustrations/.
|
|
|
|
The prompt renders each bird on a CREAM ground, not a transparent one:
|
|
the model can't cut transparency cleanly, but a flat known ground removes
|
|
cleanly in step 2. Each species gets two poses: <slug>.png (perched) and
|
|
<slug>-2.png (flight). Edit avian/scripts/prompt.template.md to change the
|
|
visual style - the prompt body is re-sent verbatim per request with
|
|
{sci_name}, {com_name}, and {pose} substituted.
|
|
|
|
Reference photos:
|
|
Cached in avian/assets/references/. The auto-fetch hits the
|
|
Wikipedia article's first image. If a reference for the species
|
|
doesn't exist locally, pregen.py fetches one and caches it. To use
|
|
a hand-picked reference, drop it in references/ named <slug>.jpg
|
|
or <slug>.png BEFORE running and pregen.py will use that instead.
|
|
|
|
Contrastive anti-reference:
|
|
For genera where Gemini's prior collapses to a more famous
|
|
lookalike, the script attaches a photo of that lookalike as a
|
|
negative reference and rewrites the prompt body to tell the model
|
|
NOT to copy the lookalike's diagnostic features. Currently wired:
|
|
Blue Jay for small blue corvids (Cyanocitta, Aphelocoma, etc.) and
|
|
Barn Swallow for other swallows (Tachycineta, Progne, etc.). The
|
|
anti-reference photos live at avian/assets/references/_anti_*.jpg
|
|
and the registry (ANTI_REFS, ANTI_REF_TRIGGERS) is keyed so adding
|
|
a new one is one entry per table.
|
|
|
|
Usage:
|
|
# Every species BirdNET-Pi knows:
|
|
python3 pregen.py --labels ~/BirdNET-Pi/model/labels.txt
|
|
|
|
# Only species observed in an eBird region:
|
|
python3 pregen.py --labels ~/BirdNET-Pi/model/labels.txt \\
|
|
--ebird-region US-CA --ebird-key YOUR_KEY
|
|
|
|
# Re-render a single species (useful after editing the prompt):
|
|
python3 pregen.py --species "Calypte anna|Anna's Hummingbird" --force
|
|
|
|
# Re-render everything after a prompt change:
|
|
python3 pregen.py --labels ~/BirdNET-Pi/model/labels.txt --force
|
|
|
|
Set GEMINI_API_KEY in the environment (preferred) or pass --gemini-key.
|
|
"""
|
|
from __future__ import annotations
|
|
import argparse
|
|
import base64
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
import time
|
|
import urllib.error
|
|
import urllib.parse
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
# Gemini's image-out model. The endpoint changes occasionally; if you
|
|
# get a 404 here, check Google's model catalog and bump this.
|
|
GEMINI_URL = (
|
|
"https://generativelanguage.googleapis.com/v1beta/models/"
|
|
"gemini-2.5-flash-image:generateContent"
|
|
)
|
|
POSES = {1: "perched", 2: "in flight with wings spread"}
|
|
|
|
# Genera where Gemini's prior collapses to Blue Jay markings unless we
|
|
# attach a Blue Jay anti-reference. Add to this set if you find another
|
|
# blue-songbird genus that needs the contrastive nudge.
|
|
JAY_GENERA = {
|
|
"Cyanocitta", "Aphelocoma", "Cyanolyca", "Calocitta", "Cyanopica",
|
|
"Garrulus", "Cyanocorax", "Gymnorhinus",
|
|
}
|
|
|
|
# Genera where Gemini's prior collapses to Barn Swallow (rufous throat,
|
|
# deeply forked tail) unless we attach a Barn Swallow anti-reference.
|
|
# Hirundo rustica is itself the Barn Swallow so it's excluded.
|
|
SWALLOW_GENERA = {
|
|
"Tachycineta", "Riparia", "Progne", "Petrochelidon", "Stelgidopteryx",
|
|
}
|
|
|
|
# Genera where Gemini's prior collapses to American Robin (gray back,
|
|
# orange breast) for ground-foraging thrushes. Add as needed.
|
|
ROBIN_GENERA = set() # placeholder for future use
|
|
|
|
# Anti-reference catalogue. Each entry describes one lookalike species
|
|
# that Gemini collapses to: the common/scientific names go in IMAGE 2's
|
|
# caption and the prompt-body bullet, and `do_not_copy` is the list of
|
|
# its diagnostic features the model must avoid. The key matches the
|
|
# `_anti_<key>.jpg` filename in the references directory and feeds into
|
|
# `ANTI_REF_TRIGGERS` below.
|
|
# ---- Style references ----
|
|
# Edo-period kachō-e woodblock prints by Ohara Koson and Hiroshi Yoshida,
|
|
# kept in a local directory (default avian/assets/references/styles/). Mapped by
|
|
# genus + pose. The bird in each print is irrelevant - only the painting
|
|
# technique (flat washes, confident outlines, tonal ground) is borrowed.
|
|
STYLE_REFS = {
|
|
"small_songbird_perched": "01-sparrows-on-bamboo-Koson.jpg",
|
|
"dark_bird_perched": "02-cawing-crow-Koson.jpg",
|
|
"vivid_perched": "03-jays-on-berry-tree-Koson.jpg",
|
|
"vibrant_perched": "04-kingfisher-Koson.jpg",
|
|
"owl": "05-owl-on-ginkgo-Koson.jpg",
|
|
"large_flight": "06-goose-flying-in-moonlight-Koson.jpg",
|
|
"small_flight": "07-swallows-in-flight-Koson.jpg",
|
|
"wader": "08-crane-in-small-water-Koson.jpg",
|
|
"pale_perched": "09-cockatoo-Yoshida.jpg",
|
|
"waterfowl_perched": "10-mandarin-ducks-Yoshida.jpg",
|
|
}
|
|
|
|
# Genus → perched style category. The first match wins. Fallback is
|
|
# "small_songbird_perched" (Koson sparrows-on-bamboo) for every uncategorized
|
|
# genus - covers passer/melospiza/spizella/junco/etc.
|
|
GENUS_STYLE_PERCHED = {
|
|
# Owls
|
|
"Tyto":"owl","Bubo":"owl","Asio":"owl","Megascops":"owl","Athene":"owl",
|
|
"Strix":"owl","Glaucidium":"owl","Aegolius":"owl",
|
|
# Hummingbirds + jays + colorful crested (vibrant color anchor)
|
|
"Calypte":"vibrant_perched","Archilochus":"vibrant_perched",
|
|
"Selasphorus":"vibrant_perched","Calothorax":"vibrant_perched",
|
|
"Cyanocitta":"vibrant_perched","Aphelocoma":"vibrant_perched",
|
|
"Pica":"vibrant_perched","Nucifraga":"vibrant_perched",
|
|
"Perisoreus":"vibrant_perched",
|
|
# Waxwings + orioles + tanagers (vivid perching with berry-tree composition feel)
|
|
"Bombycilla":"vivid_perched","Icterus":"vivid_perched",
|
|
"Piranga":"vivid_perched","Pheucticus":"vivid_perched",
|
|
"Passerina":"vivid_perched","Cardellina":"vivid_perched",
|
|
"Setophaga":"vivid_perched","Icteria":"vivid_perched",
|
|
# Corvids + vultures (dark perching)
|
|
"Corvus":"dark_bird_perched","Coragyps":"dark_bird_perched",
|
|
"Cathartes":"dark_bird_perched","Gymnogyps":"dark_bird_perched",
|
|
# Waterfowl perched (mandarin-ducks anchor)
|
|
"Anas":"waterfowl_perched","Aix":"waterfowl_perched","Mareca":"waterfowl_perched",
|
|
"Spatula":"waterfowl_perched","Branta":"waterfowl_perched","Anser":"waterfowl_perched",
|
|
"Cygnus":"waterfowl_perched","Aythya":"waterfowl_perched",
|
|
"Bucephala":"waterfowl_perched","Lophodytes":"waterfowl_perched",
|
|
"Mergus":"waterfowl_perched","Oxyura":"waterfowl_perched",
|
|
"Podiceps":"waterfowl_perched","Podilymbus":"waterfowl_perched",
|
|
"Aechmophorus":"waterfowl_perched","Gavia":"waterfowl_perched",
|
|
"Pelecanus":"waterfowl_perched","Phalacrocorax":"waterfowl_perched",
|
|
"Urile":"waterfowl_perched",
|
|
# Waders + herons (crane-in-reeds anchor)
|
|
"Ardea":"wader","Egretta":"wader","Bubulcus":"wader","Butorides":"wader",
|
|
"Nycticorax":"wader","Plegadis":"wader","Limosa":"wader","Numenius":"wader",
|
|
"Himantopus":"wader","Recurvirostra":"wader","Charadrius":"wader",
|
|
"Actitis":"wader","Calidris":"wader","Tringa":"wader",
|
|
# Pale-bodied (gulls, terns, skimmer - cockatoo anchor)
|
|
"Larus":"pale_perched","Leucophaeus":"pale_perched","Sterna":"pale_perched",
|
|
"Thalasseus":"pale_perched","Hydroprogne":"pale_perched","Rynchops":"pale_perched",
|
|
}
|
|
|
|
# Genera that should use large_flight (instead of small_flight) for pose 2.
|
|
LARGE_FLIGHT_GENERA = {
|
|
"Tyto","Bubo","Asio","Megascops","Athene","Strix","Glaucidium","Aegolius",
|
|
"Anas","Aix","Mareca","Spatula","Branta","Anser","Cygnus","Aythya",
|
|
"Bucephala","Lophodytes","Mergus","Oxyura","Pelecanus","Phalacrocorax",
|
|
"Urile","Ardea","Egretta","Bubulcus","Butorides","Nycticorax","Plegadis",
|
|
"Limosa","Numenius","Himantopus","Recurvirostra",
|
|
"Buteo","Accipiter","Aquila","Circus","Falco","Cathartes","Coragyps",
|
|
"Haliaeetus","Pandion","Elanus","Gymnogyps","Corvus",
|
|
}
|
|
|
|
|
|
def select_style_ref(sci: str, pose: int) -> str:
|
|
"""Pick the style reference filename for a (sci, pose) pair."""
|
|
genus = sci.split()[0]
|
|
# Custom overrides for hard species
|
|
if sci == "Aeronautes saxatalis":
|
|
return STYLE_REFS["vibrant_perched"] # Koson kingfisher (vertical aerial-feeder posture, no swallow bias)
|
|
if pose == 2:
|
|
return STYLE_REFS["large_flight" if genus in LARGE_FLIGHT_GENERA else "small_flight"]
|
|
return STYLE_REFS[GENUS_STYLE_PERCHED.get(genus, "small_songbird_perched")]
|
|
|
|
|
|
ANTI_REFS = {
|
|
"bluejay": {
|
|
"common_name": "Blue Jay",
|
|
"sci_name": "Cyanocitta cristata",
|
|
"do_not_copy": (
|
|
"its facial mask, its white wingbars, its black necklace, "
|
|
"its crest pattern, or its white-tipped tail"
|
|
),
|
|
},
|
|
"barnswallow": {
|
|
"common_name": "Barn Swallow",
|
|
"sci_name": "Hirundo rustica",
|
|
"do_not_copy": (
|
|
"its deep rufous throat, its long deeply forked outer tail "
|
|
"streamers, or its blue-black back"
|
|
),
|
|
},
|
|
}
|
|
|
|
# Which anti-ref to attach for which genus, and the species to exclude
|
|
# (the lookalike itself - we don't want to tell a Barn Swallow generation
|
|
# not to look like a Barn Swallow). Order matters only if a genus appears
|
|
# in more than one set; first match wins.
|
|
ANTI_REF_TRIGGERS = (
|
|
(JAY_GENERA, "bluejay", "Cyanocitta cristata"),
|
|
(SWALLOW_GENERA, "barnswallow", "Hirundo rustica"),
|
|
)
|
|
|
|
USER_AGENT = "AvianVisitors/1.0 (https://github.com/Twarner491/AvianVisitors)"
|
|
|
|
|
|
def slugify(sci: str) -> str:
|
|
"""Match avian/frontend/apt.js slugify() exactly."""
|
|
return re.sub(r"[^a-z0-9]+", "-", sci.lower()).strip("-")
|
|
|
|
|
|
def parse_species_line(line: str) -> tuple[str, str] | None:
|
|
"""Accept any of: 'Sci|Com', 'Sci_Com', 'Sci,Com'. Skip blanks + #."""
|
|
line = line.strip()
|
|
if not line or line.startswith("#"):
|
|
return None
|
|
for sep in ("|", "_", ","):
|
|
if sep in line:
|
|
sci, com = line.split(sep, 1)
|
|
sci, com = sci.strip(), com.strip()
|
|
if sci and com:
|
|
return (sci, com)
|
|
return None
|
|
|
|
|
|
def parse_species_list(lines: list[str]) -> tuple[list[tuple[str, str]], int]:
|
|
"""Returns (parsed, skipped_count)."""
|
|
out, skipped = [], 0
|
|
for line in lines:
|
|
parsed = parse_species_line(line)
|
|
if parsed:
|
|
out.append(parsed)
|
|
elif line.strip() and not line.lstrip().startswith("#"):
|
|
skipped += 1
|
|
return out, skipped
|
|
|
|
|
|
def load_prompt(path: Path) -> str:
|
|
"""Return everything after the `## Prompt` heading, stripped to the
|
|
next `##` heading (so doc preamble or trailing sections don't bleed
|
|
into the API call)."""
|
|
text = path.read_text()
|
|
m = re.search(r"##\s*Prompt\s*\n(.+?)(?=\n##\s|\Z)", text, flags=re.DOTALL)
|
|
return (m.group(1) if m else text).strip()
|
|
|
|
|
|
def ebird_filter(species, region: str, key: str):
|
|
"""Intersect a label set with the eBird species list for a region.
|
|
Region codes: US-CA (state), US-CA-085 (county)."""
|
|
url = f"https://api.ebird.org/v2/product/spplist/{region}"
|
|
req = urllib.request.Request(url, headers={"X-eBirdApiToken": key})
|
|
with urllib.request.urlopen(req, timeout=30) as r:
|
|
ebird_codes = set(json.loads(r.read()))
|
|
tax_url = "https://api.ebird.org/v2/ref/taxonomy/ebird?fmt=json"
|
|
req2 = urllib.request.Request(tax_url, headers={"X-eBirdApiToken": key})
|
|
with urllib.request.urlopen(req2, timeout=60) as r:
|
|
taxonomy = json.loads(r.read())
|
|
code_to_sci = {t["speciesCode"]: t["sciName"] for t in taxonomy}
|
|
allowed = {code_to_sci[c] for c in ebird_codes if c in code_to_sci}
|
|
return [(s, c) for s, c in species if s in allowed]
|
|
|
|
|
|
# ---- Reference photo handling ----
|
|
|
|
# Extensions to scan/write for cached reference photos. .jpg first matches
|
|
# the JPEG majority of Wikipedia infoboxes and any legacy hand-placed
|
|
# plates; .png covers PNG infoboxes (range maps, some illustrations) and
|
|
# hand-placed PNG plates. Each extension here must have a matching MIME
|
|
# in _mime_for - keep the two in sync.
|
|
REF_EXTS = (".jpg", ".png")
|
|
|
|
|
|
def fetch_wikipedia_thumb(sci: str, com: str) -> tuple[bytes, str] | None:
|
|
"""Fetch the Wikipedia article's lead/infobox image bytes.
|
|
|
|
Returns (bytes, ext) where ext is '.jpg' or '.png' sniffed from the
|
|
magic bytes - Wikipedia's infobox image can be either, and shipping
|
|
PNG bytes labeled as JPEG to Gemini gets the reference silently
|
|
rejected. Returns None if no usable image. Pi-friendly: pulls a
|
|
1024-wide thumbnail via the REST summary endpoint (a few KB to MB,
|
|
not the original-sized image).
|
|
"""
|
|
titles = [sci.replace(" ", "_"), com.replace(" ", "_"), com.split()[0]]
|
|
for title in titles:
|
|
url = (
|
|
"https://en.wikipedia.org/api/rest_v1/page/summary/"
|
|
+ urllib.parse.quote(title)
|
|
)
|
|
try:
|
|
req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
|
|
with urllib.request.urlopen(req, timeout=20) as r:
|
|
meta = json.loads(r.read())
|
|
except (urllib.error.HTTPError, urllib.error.URLError):
|
|
continue
|
|
# Prefer originalimage (higher res) over thumbnail.
|
|
for k in ("originalimage", "thumbnail"):
|
|
src = (meta.get(k) or {}).get("source")
|
|
if not src or not src.lower().endswith((".jpg", ".jpeg", ".png")):
|
|
continue
|
|
try:
|
|
req2 = urllib.request.Request(src, headers={"User-Agent": USER_AGENT})
|
|
with urllib.request.urlopen(req2, timeout=30) as r:
|
|
data = r.read()
|
|
except (urllib.error.HTTPError, urllib.error.URLError):
|
|
continue
|
|
# Magic-byte sniff - URL extension is a hint, the bytes are
|
|
# what Gemini's MIME check sees. Skip unknown formats rather
|
|
# than mis-label them.
|
|
if data.startswith(b"\x89PNG\r\n\x1a\n"):
|
|
return data, ".png"
|
|
if data.startswith(b"\xff\xd8\xff"):
|
|
return data, ".jpg"
|
|
return None
|
|
|
|
|
|
def ensure_reference(refs_dir: Path, slug: str, sci: str, com: str) -> Path | None:
|
|
"""Cache-or-fetch a reference photo. Returns the path if we have one,
|
|
None if Wikipedia had no usable image. Pre-existing references (e.g.
|
|
hand-picked Audubon plates dropped in by the user as either
|
|
<slug>.jpg or <slug>.png) are respected; the file is saved with the
|
|
extension that matches its actual format so _mime_for ships the
|
|
right MIME to Gemini."""
|
|
refs_dir.mkdir(parents=True, exist_ok=True)
|
|
for ext in REF_EXTS:
|
|
cached = refs_dir / f"{slug}{ext}"
|
|
if cached.exists() and cached.stat().st_size > 1024:
|
|
return cached
|
|
fetched = fetch_wikipedia_thumb(sci, com)
|
|
if not fetched:
|
|
return None
|
|
data, ext = fetched
|
|
path = refs_dir / f"{slug}{ext}"
|
|
path.write_bytes(data)
|
|
return path
|
|
|
|
|
|
def select_anti_ref_key(sci: str) -> str | None:
|
|
"""Return the ANTI_REFS key for the lookalike that Gemini drifts
|
|
toward for this species, or None if no anti-ref is needed. The key
|
|
matches `_anti_<key>.jpg` in the references directory."""
|
|
genus = sci.split()[0]
|
|
for genera, key, exclude in ANTI_REF_TRIGGERS:
|
|
if genus in genera and sci != exclude:
|
|
return key
|
|
return None
|
|
|
|
|
|
def load_species_notes(notes_path: Path) -> dict[str, str]:
|
|
"""Load per-species prompt addenda. Keys are scientific names; values
|
|
are 1-2 sentence clarifications to inject when generating that
|
|
species. Returns {} if the notes file doesn't exist."""
|
|
if not notes_path.exists():
|
|
return {}
|
|
raw = json.loads(notes_path.read_text())
|
|
return {k: v for k, v in raw.items()
|
|
if not k.startswith("_") and isinstance(v, str)}
|
|
|
|
|
|
def load_anti_ref(refs_dir: Path, key: str = "bluejay") -> Path | None:
|
|
"""Return path to the bundled anti-reference for the given key,
|
|
if present. Known keys: bluejay, barnswallow."""
|
|
p = refs_dir / f"_anti_{key}.jpg"
|
|
return p if p.exists() else None
|
|
|
|
|
|
# ---- Gemini call ----
|
|
|
|
def _anti_ref_line(anti_ref_key: str | None) -> str:
|
|
"""Render the `{anti_ref_line}` substitution for the prompt body.
|
|
Returns the IMAGE 2 bullet describing which species is attached and
|
|
which of its features the model must avoid - or an empty string
|
|
when no anti-ref is attached for this species."""
|
|
info = ANTI_REFS.get(anti_ref_key or "")
|
|
if not info:
|
|
return ""
|
|
return (
|
|
f"- IMAGE 2 (negative, when attached) is a {info['common_name']} "
|
|
f"({info['sci_name']}). It is NOT what you are drawing. Do NOT "
|
|
f"copy {info['do_not_copy']}. If your output looks more like "
|
|
f"IMAGE 2 than IMAGE 1, the output is wrong."
|
|
)
|
|
|
|
|
|
def gen_one(
|
|
api_key: str,
|
|
prompt: str,
|
|
sci: str,
|
|
com: str,
|
|
pose: int,
|
|
positive_ref: Path | None = None,
|
|
anti_ref: Path | None = None,
|
|
anti_ref_key: str | None = None,
|
|
species_note: str | None = None,
|
|
style_ref: Path | None = None,
|
|
) -> bytes:
|
|
"""Single Gemini call with bounded retry on 429 + transient 5xx.
|
|
Returns raw PNG bytes.
|
|
|
|
positive_ref: Wikipedia/Audubon photo of the target species.
|
|
anti_ref: lookalike photo to attach as IMAGE 2. The companion
|
|
anti_ref_key (a key into ANTI_REFS) must match what's in
|
|
the file - it drives the IMAGE 2 caption and the
|
|
{anti_ref_line} substitution in the prompt body. Pass
|
|
both or neither; passing the path without the key would
|
|
caption the image as an unnamed "another species".
|
|
species_note: optional 1-2 sentence clarifier for difficult species,
|
|
appended as the last paragraph before the reference
|
|
block.
|
|
"""
|
|
body = (prompt
|
|
.replace("{sci_name}", sci)
|
|
.replace("{com_name}", com)
|
|
.replace("{pose}", POSES[pose])
|
|
.replace("{anti_ref_line}", _anti_ref_line(anti_ref_key)))
|
|
if species_note:
|
|
body = body + "\n\nSpecies-specific note: " + species_note
|
|
|
|
parts: list[dict] = [{"text": body}]
|
|
if positive_ref:
|
|
# Downscale the anatomy reference to 384px on the long side
|
|
# before encoding. Big Wikipedia photos visually dominate as a
|
|
# style signal even though the prompt says they're anatomy-only;
|
|
# at 384px the model still reads species/markings/colors but
|
|
# has less photographic detail to mimic.
|
|
try:
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
img = Image.open(positive_ref).convert("RGB")
|
|
w, h = img.size
|
|
if max(w, h) > 384:
|
|
scale = 384 / max(w, h)
|
|
img = img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
|
|
buf = BytesIO()
|
|
img.save(buf, format="PNG", optimize=True)
|
|
ref_bytes = buf.getvalue()
|
|
ref_mime = "image/png"
|
|
except Exception:
|
|
ref_bytes = positive_ref.read_bytes()
|
|
ref_mime = _mime_for(positive_ref)
|
|
parts.append({"text": "IMAGE 1 (positive, target species):"})
|
|
parts.append({"inline_data": {
|
|
"mime_type": ref_mime,
|
|
"data": base64.b64encode(ref_bytes).decode(),
|
|
}})
|
|
if anti_ref:
|
|
anti_name = (ANTI_REFS.get(anti_ref_key or "") or {}).get(
|
|
"common_name", "lookalike species"
|
|
)
|
|
parts.append({"text": f"IMAGE 2 (negative, {anti_name}, do NOT copy):"})
|
|
parts.append({"inline_data": {
|
|
"mime_type": _mime_for(anti_ref),
|
|
"data": base64.b64encode(anti_ref.read_bytes()).decode(),
|
|
}})
|
|
if style_ref:
|
|
parts.append({"text": (
|
|
"IMAGE 3 (positive STYLE reference - Edo-period kachō-e woodblock "
|
|
"print). The species in IMAGE 3 is irrelevant; only its painting "
|
|
"technique is borrowed (flat washes, confident outlines, tonal "
|
|
"mineral-pigment ground). DO NOT copy any branches, leaves, water, "
|
|
"moon, or scenery from IMAGE 3.")})
|
|
parts.append({"inline_data": {
|
|
"mime_type": _mime_for(style_ref),
|
|
"data": base64.b64encode(style_ref.read_bytes()).decode(),
|
|
}})
|
|
|
|
payload = {
|
|
"contents": [{"parts": parts}],
|
|
# TEXT included so Gemini can surface safety messaging without
|
|
# rejecting the request shape (image-only sometimes errors).
|
|
"generationConfig": {"responseModalities": ["TEXT", "IMAGE"]},
|
|
}
|
|
# API key as header, NOT URL - keeps the key out of Google's
|
|
# request logs, proxy logs, and shell history.
|
|
req = urllib.request.Request(
|
|
GEMINI_URL,
|
|
data=json.dumps(payload).encode(),
|
|
headers={"Content-Type": "application/json", "x-goog-api-key": api_key},
|
|
method="POST",
|
|
)
|
|
|
|
backoff = 4.0
|
|
for attempt in range(4):
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=180) as r:
|
|
resp = json.loads(r.read())
|
|
break
|
|
except urllib.error.HTTPError as e:
|
|
if e.code in (429, 500, 502, 503, 504) and attempt < 3:
|
|
ra = e.headers.get("Retry-After")
|
|
try:
|
|
retry_after = float(ra) if ra else backoff
|
|
except (TypeError, ValueError):
|
|
retry_after = backoff # HTTP-date format, fall back
|
|
time.sleep(retry_after)
|
|
backoff *= 2
|
|
continue
|
|
raise
|
|
except urllib.error.URLError:
|
|
if attempt < 3:
|
|
time.sleep(backoff)
|
|
backoff *= 2
|
|
continue
|
|
raise
|
|
|
|
for cand in resp.get("candidates", []):
|
|
for part in cand.get("content", {}).get("parts", []):
|
|
inline = part.get("inlineData") or part.get("inline_data")
|
|
if inline and inline.get("data"):
|
|
return base64.b64decode(inline["data"])
|
|
# No image - surface the blocking reason so users know what to fix.
|
|
finish = (resp.get("candidates", [{}])[0]).get("finishReason", "?")
|
|
block = resp.get("promptFeedback", {}).get("blockReason", "")
|
|
raise RuntimeError(f"no image (finish={finish} block={block})")
|
|
|
|
|
|
def _mime_for(p: Path) -> str:
|
|
ext = p.suffix.lower()
|
|
if ext in (".jpg", ".jpeg"):
|
|
return "image/jpeg"
|
|
if ext == ".png":
|
|
return "image/png"
|
|
if ext == ".webp":
|
|
return "image/webp"
|
|
return "application/octet-stream"
|
|
|
|
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser(
|
|
description=__doc__,
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
)
|
|
src = ap.add_mutually_exclusive_group(required=True)
|
|
src.add_argument("--labels", type=Path, help="Path to BirdNET-Pi labels.txt (or any file of Sci|Com lines)")
|
|
src.add_argument("--species", action="append", default=[],
|
|
help="Manual 'Sci|Com' (repeatable)")
|
|
src.add_argument("--stdin", action="store_true", help="Read Sci|Com lines from stdin")
|
|
ap.add_argument("--ebird-region", help="eBird region code (e.g. US-CA, US-CA-085) to filter labels")
|
|
ap.add_argument("--ebird-key", help="eBird API key (or EBIRD_API_KEY env)")
|
|
ap.add_argument("--gemini-key", help="Gemini API key (or GEMINI_API_KEY env)")
|
|
ap.add_argument("--out", type=Path,
|
|
default=Path(__file__).resolve().parents[1] / "assets" / "illustrations",
|
|
help="Output directory (default: avian/assets/illustrations/)")
|
|
ap.add_argument("--refs", type=Path,
|
|
default=Path(__file__).resolve().parents[1] / "assets" / "references",
|
|
help="Reference photo cache directory (default: avian/assets/references/)")
|
|
ap.add_argument("--styles", type=Path,
|
|
default=Path(__file__).resolve().parents[1] / "assets" / "references" / "styles",
|
|
help="Style reference directory (default: avian/assets/references/styles/)")
|
|
ap.add_argument("--prompt", type=Path,
|
|
default=Path(__file__).resolve().parent / "prompt.template.md",
|
|
help="Prompt template path")
|
|
ap.add_argument("--notes", type=Path,
|
|
default=Path(__file__).resolve().parent / "species-notes.json",
|
|
help="Per-species prompt addenda for difficult cases (e.g. similar-species drift)")
|
|
ap.add_argument("--poses", nargs="+", type=int, default=[1, 2],
|
|
choices=list(POSES.keys()),
|
|
help="Which poses to render. 1=perched, 2=flight. Default: both.")
|
|
ap.add_argument("--force", action="store_true", help="Re-render even if file exists")
|
|
ap.add_argument("--no-refs", action="store_true",
|
|
help="Skip the Wikipedia reference fetch (faster, lower-quality output)")
|
|
ap.add_argument("--sleep", type=float, default=6.0,
|
|
help="Seconds between API calls (default 6 = headroom under free-tier RPM cap)")
|
|
ap.add_argument("--limit", type=int, default=0, help="Cap species count for testing")
|
|
args = ap.parse_args()
|
|
|
|
gemini_key = args.gemini_key or os.environ.get("GEMINI_API_KEY", "")
|
|
if not gemini_key:
|
|
print("error: GEMINI_API_KEY required (--gemini-key or env)", file=sys.stderr)
|
|
return 2
|
|
|
|
# Build species list
|
|
if args.labels:
|
|
species, skipped = parse_species_list(args.labels.read_text().splitlines())
|
|
elif args.stdin:
|
|
species, skipped = parse_species_list(sys.stdin.read().splitlines())
|
|
else:
|
|
species, skipped = parse_species_list(args.species)
|
|
if skipped:
|
|
print(f"[parse] skipped {skipped} malformed line(s)", file=sys.stderr)
|
|
if not species:
|
|
print("error: no species resolved", file=sys.stderr)
|
|
return 2
|
|
|
|
if args.ebird_region:
|
|
ek = args.ebird_key or os.environ.get("EBIRD_API_KEY", "")
|
|
if not ek:
|
|
print("error: --ebird-region requires --ebird-key or EBIRD_API_KEY", file=sys.stderr)
|
|
return 2
|
|
print(f"[ebird] filtering {len(species)} species against {args.ebird_region}...")
|
|
species = ebird_filter(species, args.ebird_region, ek)
|
|
|
|
if args.limit:
|
|
species = species[:args.limit]
|
|
|
|
prompt = load_prompt(args.prompt)
|
|
args.out.mkdir(parents=True, exist_ok=True)
|
|
anti_paths: dict[str, Path] = {}
|
|
if not args.no_refs:
|
|
for key in ANTI_REFS:
|
|
p = load_anti_ref(args.refs, key)
|
|
if p:
|
|
anti_paths[key] = p
|
|
notes = load_species_notes(args.notes)
|
|
if notes:
|
|
print(f"[notes] loaded per-species addenda for {len(notes)} species")
|
|
|
|
total = len(species) * len(args.poses)
|
|
print(f"generating up to {total} illustrations into {args.out}/")
|
|
for key, p in anti_paths.items():
|
|
print(f"[refs] {ANTI_REFS[key]['common_name']} anti-reference: {p.name}")
|
|
|
|
done = skipped_existing = failed = 0
|
|
first_fail = None
|
|
for idx, (sci, com) in enumerate(species):
|
|
slug = slugify(sci)
|
|
pos_ref = None
|
|
if not args.no_refs:
|
|
pos_ref = ensure_reference(args.refs, slug, sci, com)
|
|
if not pos_ref:
|
|
print(f" [warn] no Wikipedia photo for {sci} - proceeding without positive ref", file=sys.stderr)
|
|
anti_key = select_anti_ref_key(sci)
|
|
anti = anti_paths.get(anti_key) if anti_key else None
|
|
# Pair the key with the path so gen_one captions IMAGE 2 with the
|
|
# right species (the bug this fixes: stale Blue-Jay caption on
|
|
# an attached Barn-Swallow anti-ref).
|
|
anti_key_for_call = anti_key if anti else None
|
|
|
|
for pose in args.poses:
|
|
fname = f"{slug}.png" if pose == 1 else f"{slug}-{pose}.png"
|
|
path = args.out / fname
|
|
if path.exists() and not args.force:
|
|
skipped_existing += 1
|
|
continue
|
|
try:
|
|
style_ref_path = args.styles / select_style_ref(sci, pose)
|
|
if not style_ref_path.exists():
|
|
style_ref_path = None
|
|
data = gen_one(gemini_key, prompt, sci, com, pose,
|
|
positive_ref=pos_ref, anti_ref=anti,
|
|
anti_ref_key=anti_key_for_call,
|
|
species_note=notes.get(sci),
|
|
style_ref=style_ref_path)
|
|
path.write_bytes(data)
|
|
done += 1
|
|
refs_tag = "+ref" if pos_ref else ""
|
|
anti_tag = "+anti" if anti else ""
|
|
note_tag = "+note" if notes.get(sci) else ""
|
|
print(f" [ok] {fname} ({len(data)//1024} KB){refs_tag}{anti_tag}{note_tag}")
|
|
except (urllib.error.HTTPError, urllib.error.URLError, RuntimeError) as e:
|
|
failed += 1
|
|
first_fail = first_fail or fname
|
|
print(f" [fail] {fname}: {e}", file=sys.stderr)
|
|
# Don't sleep after the last species' last pose.
|
|
if not (idx == len(species) - 1 and pose == args.poses[-1]):
|
|
time.sleep(args.sleep)
|
|
|
|
print(f"\ngenerated {done} · skipped {skipped_existing} · failed {failed}")
|
|
if first_fail:
|
|
print(f"first failure: {first_fail} (re-run without --force to retry only the misses)", file=sys.stderr)
|
|
return 0 if failed == 0 else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|