mirror of
https://github.com/ostris/ai-toolkit.git
synced 2026-07-11 09:42:06 +00:00
320 lines
11 KiB
Python
320 lines
11 KiB
Python
"""Shared helpers for Ideogram-4 structured JSON captions.
|
|
|
|
This is the single source of truth for the caption schema so the captioner, the
|
|
prompt upsampler, the dataloader, and the model encoder all agree. It encodes the
|
|
official Ideogram-4 rules and, crucially, MIGRATES the old caption format we used
|
|
before those rules were published into the new one ("digest" old, emit new).
|
|
|
|
Official schema (summary):
|
|
- three top-level keys: high_level_description (optional), style_description
|
|
(optional), compositional_deconstruction (required).
|
|
- style_description holds EXACTLY ONE of `photo` (photographs) or `art_style`
|
|
(illustration/painting/3D/graphic design), never both. Key order is strict and
|
|
branch-dependent:
|
|
photo branch: aesthetics, lighting, photo, medium, color_palette
|
|
non-photo branch: aesthetics, lighting, medium, art_style, color_palette
|
|
- medium is one of: photograph, illustration, 3d_render, painting, graphic_design
|
|
- color_palette: UPPERCASE #RRGGBB only, up to 16 per image / 5 per element.
|
|
- elements, strict key order:
|
|
obj: type, bbox, desc, color_palette
|
|
text: type, bbox, text, desc, color_palette
|
|
bbox is optional, normalized 0-1000, [y_min, x_min, y_max, x_max], top-left.
|
|
- serialize compact: separators=(",", ":"), ensure_ascii=False (no \\uXXXX).
|
|
|
|
The OLD format we previously emitted differed by: always using `photo` (even for
|
|
non-photo media), putting `color_palette` before `desc`/`text`, title-cased medium
|
|
with a trailing period ("Illustration."), and lowercase / 3-digit hex. Every
|
|
function here accepts the old shape and returns the new one.
|
|
"""
|
|
|
|
import json
|
|
import re
|
|
from collections import OrderedDict
|
|
|
|
MAX_IMAGE_PALETTE = 16 # style_description.color_palette
|
|
MAX_ELEMENT_PALETTE = 5 # per-element color_palette
|
|
|
|
# Canonical medium tokens (official set).
|
|
MEDIUM_OPTIONS = [
|
|
"photograph",
|
|
"illustration",
|
|
"3d_render",
|
|
"painting",
|
|
"graphic_design",
|
|
]
|
|
|
|
# Map common variants (including our old "Title." style) to the canonical token.
|
|
# Anything not listed is treated as a custom medium and preserved verbatim.
|
|
_MEDIUM_ALIASES = {
|
|
"photograph": "photograph",
|
|
"photo": "photograph",
|
|
"illustration": "illustration",
|
|
"3d render": "3d_render",
|
|
"3d_render": "3d_render",
|
|
"3d-render": "3d_render",
|
|
"3drender": "3d_render",
|
|
"render": "3d_render",
|
|
"3d": "3d_render",
|
|
"painting": "painting",
|
|
"graphic design": "graphic_design",
|
|
"graphic_design": "graphic_design",
|
|
"graphic-design": "graphic_design",
|
|
"graphic": "graphic_design",
|
|
}
|
|
|
|
_HEX6_RE = re.compile(r"^#[0-9a-fA-F]{6}$")
|
|
_HEX3_RE = re.compile(r"^#[0-9a-fA-F]{3}$")
|
|
|
|
|
|
def canon_medium(medium):
|
|
"""Canonicalize a medium string to an official token when recognized,
|
|
otherwise return it stripped (custom mediums are allowed, preserved as-is)."""
|
|
if not isinstance(medium, str):
|
|
return medium
|
|
key = medium.strip().rstrip(".").strip().lower()
|
|
if key in _MEDIUM_ALIASES:
|
|
return _MEDIUM_ALIASES[key]
|
|
return medium.strip()
|
|
|
|
|
|
def is_photo_medium(medium):
|
|
"""True for the photograph branch (uses `photo`), False for the art_style branch."""
|
|
return canon_medium(medium) == "photograph"
|
|
|
|
|
|
def normalize_hex(color):
|
|
"""Return an UPPERCASE #RRGGBB string, expanding #RGB -> #RRGGBB. None if invalid."""
|
|
if not isinstance(color, str):
|
|
return None
|
|
s = color.strip()
|
|
if _HEX6_RE.match(s):
|
|
return "#" + s[1:].upper()
|
|
if _HEX3_RE.match(s):
|
|
return "#" + "".join(ch * 2 for ch in s[1:]).upper()
|
|
return None
|
|
|
|
|
|
def sanitize_palette(palette, max_len):
|
|
"""Keep unique, valid, UPPERCASE hex colors in order, capped to max_len.
|
|
Returns the cleaned list, or None if nothing valid remains (drop the key)."""
|
|
if not isinstance(palette, (list, tuple)):
|
|
return None
|
|
seen = set()
|
|
out = []
|
|
for c in palette:
|
|
h = normalize_hex(c)
|
|
if h is None or h in seen:
|
|
continue
|
|
seen.add(h)
|
|
out.append(h)
|
|
if len(out) >= max_len:
|
|
break
|
|
return out or None
|
|
|
|
|
|
def normalize_style(style):
|
|
"""Reorder/clean style_description into the correct branch (photo vs art_style)
|
|
with the strict key order, canonical medium, and uppercase palette. Accepts the
|
|
old shape (always `photo`) and migrates it based on the medium."""
|
|
if not isinstance(style, dict):
|
|
return style
|
|
|
|
raw_medium = style.get("medium")
|
|
medium = canon_medium(raw_medium) if raw_medium is not None else None
|
|
has_photo = bool(style.get("photo"))
|
|
has_art = bool(style.get("art_style"))
|
|
|
|
# Decide the branch. A recognized medium is authoritative; otherwise infer from
|
|
# whichever render key the (old) data already had, defaulting to photo.
|
|
if medium in MEDIUM_OPTIONS:
|
|
photo_branch = medium == "photograph"
|
|
elif has_art and not has_photo:
|
|
photo_branch = False
|
|
else:
|
|
photo_branch = True
|
|
|
|
photo_val = style.get("photo") if has_photo else None
|
|
art_val = style.get("art_style") if has_art else None
|
|
|
|
out = OrderedDict()
|
|
if "aesthetics" in style:
|
|
out["aesthetics"] = style["aesthetics"]
|
|
if "lighting" in style:
|
|
out["lighting"] = style["lighting"]
|
|
|
|
if photo_branch:
|
|
# aesthetics, lighting, photo, medium, color_palette
|
|
val = photo_val if photo_val is not None else art_val
|
|
if val is not None:
|
|
out["photo"] = val
|
|
if medium is not None:
|
|
out["medium"] = medium
|
|
else:
|
|
# aesthetics, lighting, medium, art_style, color_palette
|
|
if medium is not None:
|
|
out["medium"] = medium
|
|
val = art_val if art_val is not None else photo_val
|
|
if val is not None:
|
|
out["art_style"] = val
|
|
|
|
pal = sanitize_palette(style.get("color_palette"), MAX_IMAGE_PALETTE)
|
|
if pal is not None:
|
|
out["color_palette"] = pal
|
|
|
|
# Preserve any unexpected extra keys at the end rather than dropping them.
|
|
for k, v in style.items():
|
|
if k not in (
|
|
"aesthetics",
|
|
"lighting",
|
|
"photo",
|
|
"art_style",
|
|
"medium",
|
|
"color_palette",
|
|
):
|
|
out[k] = v
|
|
return out
|
|
|
|
|
|
def normalize_element(el):
|
|
"""Reorder an element's keys to the strict schema order and uppercase its
|
|
palette. obj: type, bbox, desc, color_palette. text: type, bbox, text, desc,
|
|
color_palette. bbox is kept verbatim (already [y1,x1,y2,x2] in stored form)."""
|
|
if not isinstance(el, dict):
|
|
return el
|
|
etype = el.get("type", "obj")
|
|
out = OrderedDict()
|
|
out["type"] = etype
|
|
if el.get("bbox") is not None:
|
|
out["bbox"] = el["bbox"]
|
|
if etype == "text":
|
|
if "text" in el:
|
|
out["text"] = el["text"]
|
|
if "desc" in el:
|
|
out["desc"] = el["desc"]
|
|
else:
|
|
if "desc" in el:
|
|
out["desc"] = el["desc"]
|
|
pal = sanitize_palette(el.get("color_palette"), MAX_ELEMENT_PALETTE)
|
|
if pal is not None:
|
|
out["color_palette"] = pal
|
|
# Preserve any extras (e.g. future keys) at the end.
|
|
for k, v in el.items():
|
|
if k not in out and k != "color_palette":
|
|
out[k] = v
|
|
return out
|
|
|
|
|
|
def normalize_caption_dict(data):
|
|
"""Normalize a parsed caption dict in place-ish: drop input-only aspect_ratio,
|
|
enforce top-level key order, normalize style (photo/art_style branch) and every
|
|
element. Returns a new OrderedDict. Accepts old-format captions and emits new."""
|
|
if not isinstance(data, dict):
|
|
return data
|
|
data.pop("aspect_ratio", None) # input-only context, never part of output
|
|
|
|
out = OrderedDict()
|
|
if "high_level_description" in data:
|
|
out["high_level_description"] = data["high_level_description"]
|
|
if "style_description" in data:
|
|
out["style_description"] = normalize_style(data["style_description"])
|
|
|
|
decon = data.get("compositional_deconstruction")
|
|
if isinstance(decon, dict):
|
|
nd = OrderedDict()
|
|
if "background" in decon:
|
|
nd["background"] = decon["background"]
|
|
els = decon.get("elements")
|
|
if isinstance(els, list):
|
|
nd["elements"] = [normalize_element(e) for e in els]
|
|
for k, v in decon.items():
|
|
if k not in ("background", "elements"):
|
|
nd[k] = v
|
|
out["compositional_deconstruction"] = nd
|
|
elif decon is not None:
|
|
out["compositional_deconstruction"] = decon
|
|
|
|
for k, v in data.items():
|
|
if k not in (
|
|
"high_level_description",
|
|
"style_description",
|
|
"compositional_deconstruction",
|
|
"aspect_ratio",
|
|
):
|
|
out[k] = v
|
|
return out
|
|
|
|
|
|
# --- bbox coordinate adaptation that does NOT require valid JSON -------------
|
|
# Captioners emit boxes as [x1,y1,x2,y2] but we store [y1,x1,y2,x2]. The
|
|
# structured normalizer can only swap per-element when the JSON parses; if the
|
|
# model returns malformed JSON, that path is skipped and the boxes stay in the
|
|
# wrong order. This regex rewrites every `"bbox":[...]` array in the raw text
|
|
# directly, so the swap still happens on un-parseable output.
|
|
_BBOX_TEXT_RE = re.compile(
|
|
r'"bbox"\s*:\s*\[\s*'
|
|
r"(-?\d+(?:\.\d+)?)\s*,\s*"
|
|
r"(-?\d+(?:\.\d+)?)\s*,\s*"
|
|
r"(-?\d+(?:\.\d+)?)\s*,\s*"
|
|
r"(-?\d+(?:\.\d+)?)\s*\]"
|
|
)
|
|
|
|
|
|
def _clamp_1000(v):
|
|
return max(0, min(1000, round(float(v))))
|
|
|
|
|
|
def swap_bbox_xy_in_text(text):
|
|
"""Swap every [x1,y1,x2,y2] bbox to the stored [y1,x1,y2,x2] order directly in
|
|
the raw model output -- clamping each value to 0-1000 and ordering each axis
|
|
pair. It never parses the surrounding JSON, so it works even when the output is
|
|
malformed. Only `"bbox":[n,n,n,n]` arrays are touched; everything else is left
|
|
byte-for-byte. Returns the rewritten text."""
|
|
if not isinstance(text, str):
|
|
return text
|
|
|
|
def _repl(m):
|
|
x1, y1, x2, y2 = m.groups()
|
|
cx1, cx2 = sorted((_clamp_1000(x1), _clamp_1000(x2)))
|
|
cy1, cy2 = sorted((_clamp_1000(y1), _clamp_1000(y2)))
|
|
return f'"bbox":[{cy1},{cx1},{cy2},{cx2}]'
|
|
|
|
return _BBOX_TEXT_RE.sub(_repl, text)
|
|
|
|
|
|
def is_ideogram_caption_str(text):
|
|
"""True if text parses as a JSON object with a compositional_deconstruction block."""
|
|
t = (text or "").strip()
|
|
if not t.startswith("{"):
|
|
return False
|
|
try:
|
|
d = json.loads(t)
|
|
except Exception:
|
|
return False
|
|
return isinstance(d, dict) and isinstance(
|
|
d.get("compositional_deconstruction"), dict
|
|
)
|
|
|
|
|
|
def to_model_string(data):
|
|
"""Serialize a caption dict to the compact, model-ready string the renderer wants."""
|
|
return json.dumps(data, ensure_ascii=False, separators=(",", ":"))
|
|
|
|
|
|
def digest_caption_string(text):
|
|
"""Parse, normalize (migrating old format), and return the compact model-ready
|
|
string. Returns the input unchanged if it is not an Ideogram structured caption
|
|
(plain-text captions pass straight through)."""
|
|
t = (text or "").strip()
|
|
if not t.startswith("{"):
|
|
return text
|
|
try:
|
|
data = json.loads(t, object_pairs_hook=OrderedDict)
|
|
except Exception:
|
|
return text
|
|
if not (
|
|
isinstance(data, dict)
|
|
and isinstance(data.get("compositional_deconstruction"), dict)
|
|
):
|
|
return text
|
|
return to_model_string(normalize_caption_dict(data))
|