Spatial9/GravityLLM
0
1from __future__ import annotations2 3import json4import math5import re6from pathlib import Path7from typing import Any, Dict, List, Tuple8 9import matplotlib.pyplot as plt10from jsonschema import Draft7Validator11 12 13SCHEMA_PATH = Path(__file__).resolve().parents[1] / "schemas" / "scene.schema.json"14SCHEMA: Dict[str, Any] = json.loads(SCHEMA_PATH.read_text(encoding="utf-8"))15VALIDATOR = Draft7Validator(SCHEMA)16 17CLASS_DEFAULTS = {18 "lead_vocal": {"az": 0, "el": 10, "dist": 1.6, "width": 0.15, "gain": 0.0, "rev": 0.18, "er": 0.22},19 "dialogue": {"az": 0, "el": 5, "dist": 1.5, "width": 0.10, "gain": 0.0, "rev": 0.08, "er": 0.18},20 "back_vocal": {"az": -18, "el": 8, "dist": 1.9, "width": 0.25, "gain": -1.2, "rev": 0.16, "er": 0.16},21 "kick": {"az": 0, "el": 0, "dist": 2.2, "width": 0.0, "gain": 0.0, "rev": 0.02, "er": 0.05},22 "snare": {"az": 8, "el": 4, "dist": 2.4, "width": 0.1, "gain": -0.3, "rev": 0.05, "er": 0.07},23 "hihat": {"az": 22, "el": 7, "dist": 2.7, "width": 0.18, "gain": -1.0, "rev": 0.06, "er": 0.08},24 "bass": {"az": 0, "el": -5, "dist": 2.6, "width": 0.05, "gain": -0.5, "rev": 0.03, "er": 0.06},25 "drums_bus": {"az": 0, "el": 2, "dist": 2.8, "width": 0.25, "gain": -0.6, "rev": 0.08, "er": 0.10},26 "pad": {"az": -70, "el": 18, "dist": 4.8, "width": 0.82, "gain": -4.0, "rev": 0.28, "er": 0.12},27 "synth_lead": {"az": 24, "el": 15, "dist": 2.0, "width": 0.35, "gain": -1.0, "rev": 0.12, "er": 0.10},28 "guitar": {"az": -28, "el": 10, "dist": 2.8, "width": 0.28, "gain": -1.5, "rev": 0.09, "er": 0.11},29 "piano": {"az": -14, "el": 8, "dist": 2.6, "width": 0.34, "gain": -1.5, "rev": 0.09, "er": 0.10},30 "fx": {"az": 105, "el": 28, "dist": 6.2, "width": 0.66, "gain": -7.0, "rev": 0.42, "er": 0.08},31 "ambience": {"az": -120, "el": 15, "dist": 8.0, "width": 0.90, "gain": -8.0, "rev": 0.55, "er": 0.10},32 "other": {"az": 0, "el": 8, "dist": 3.0, "width": 0.25, "gain": -2.0, "rev": 0.10, "er": 0.10},33}34 35SPATIAL_CLASSES = list(CLASS_DEFAULTS)36 37 38def clip(value: float, lo: float, hi: float) -> float:39 return max(lo, min(hi, value))40 41 42def extract_first_json_block(text: str) -> str:43 match = re.search(r"\{.*\}", text, flags=re.DOTALL)44 return match.group(0).strip() if match else text.strip()45 46 47def parse_json_text(text: str) -> Dict[str, Any]:48 return json.loads(extract_first_json_block(text))49 50 51def validate_scene(scene: Dict[str, Any]) -> Tuple[bool, List[str]]:52 errors = sorted(VALIDATOR.iter_errors(scene), key=lambda e: list(e.path))53 if not errors:54 return True, []55 messages = []56 for err in errors[:50]:57 path = ".".join(str(x) for x in err.path)58 messages.append(f"{path or '<root>'}: {err.message}")59 return False, messages60 61 62def make_motion(az: float, el: float, dist: float, sweep: float = 0.0) -> List[Dict[str, float]]:63 if abs(sweep) < 0.001:64 return [65 {"t": 0.0, "az_deg": round(az, 2), "el_deg": round(el, 2), "dist_m": round(dist, 2)},66 {"t": 1.0, "az_deg": round(az, 2), "el_deg": round(el, 2), "dist_m": round(dist, 2)},67 ]68 return [69 {"t": 0.0, "az_deg": round(az - sweep / 2, 2), "el_deg": round(el, 2), "dist_m": round(dist, 2)},70 {"t": 1.0, "az_deg": round(az + sweep / 2, 2), "el_deg": round(el, 2), "dist_m": round(dist, 2)},71 ]72 73 74def find_anchor_rules(payload: Dict[str, Any]) -> Dict[str, Dict[str, float]]:75 anchors: Dict[str, Dict[str, float]] = {}76 for rule in payload.get("rules", []):77 if rule.get("type") == "anchor":78 key = rule.get("track_class")79 if key:80 anchors[key] = {81 "az": float(rule.get("az_deg", 0.0)),82 "el": float(rule.get("el_deg", 0.0)),83 "dist": float(rule.get("dist_m", 1.6)),84 }85 return anchors86 87 88def width_preference(payload: Dict[str, Any], track_class: str) -> float | None:89 for rule in payload.get("rules", []):90 if rule.get("type") == "width_pref" and rule.get("track_class") == track_class:91 value = rule.get("min_width")92 if value is not None:93 return float(value)94 return None95 96 97def target_layout(payload: Dict[str, Any]) -> str:98 fmt = str(payload.get("target_format", "iamf")).lower()99 if fmt in {"binaural", "iamf", "5.1.4", "7.1.4"}:100 return fmt101 return "iamf"102 103 104def room_preset(payload: Dict[str, Any]) -> str:105 style = str(payload.get("style", "neutral")).lower()106 mapping = {107 "club": "club_medium",108 "cinematic": "cinema_large",109 "film": "cinema_large",110 "podcast": "studio_dry",111 "live": "stage_wide",112 "intimate": "studio_small",113 }114 return mapping.get(style, "studio_neutral")115 116 117def heuristic_scene(payload: Dict[str, Any]) -> Dict[str, Any]:118 anchors = find_anchor_rules(payload)119 objects: List[Dict[str, Any]] = []120 constraints_applied: List[str] = []121 max_objects = int(payload.get("max_objects", 10))122 style = str(payload.get("style", "neutral")).lower()123 124 for rule in payload.get("rules", []):125 rtype = rule.get("type")126 if rtype == "anchor":127 constraints_applied.append(128 f"anchor:{rule.get('track_class')}@{rule.get('az_deg')}/{rule.get('el_deg')}/{rule.get('dist_m')}"129 )130 elif rtype == "mono_low_end":131 constraints_applied.append(f"mono_low_end<{rule.get('hz_below', 120)}Hz")132 elif rtype == "width_pref":133 constraints_applied.append(f"{rule.get('track_class')}_width>={rule.get('min_width')}")134 elif rtype == "keep_dialogue_clear":135 band = rule.get("band_hz", [1000, 4000])136 constraints_applied.append(f"keep_dialogue_clear_{band[0]}-{band[1]}Hz")137 elif rtype == "avoid_band_masking":138 band = rule.get("band_hz", [1500, 4500])139 constraints_applied.append(f"avoid_masking_{rule.get('mask_target')}_{band[0]}-{band[1]}Hz")140 141 for idx, stem in enumerate(payload.get("stems", [])[:max_objects]):142 cls = stem.get("class", "other")143 defaults = dict(CLASS_DEFAULTS.get(cls, CLASS_DEFAULTS["other"]))144 if cls in anchors:145 defaults["az"] = anchors[cls]["az"]146 defaults["el"] = anchors[cls]["el"]147 defaults["dist"] = anchors[cls]["dist"]148 149 width_min = width_preference(payload, cls)150 if width_min is not None:151 defaults["width"] = max(defaults["width"], width_min)152 153 leadness = float(stem.get("leadness", 0.0))154 transient = float(stem.get("transient", 0.0))155 lufs = float(stem.get("lufs", -20.0))156 157 if leadness > 0.8 and cls not in {"kick", "bass", "dialogue", "lead_vocal"}:158 defaults["dist"] = clip(defaults["dist"] - 0.35, 0.8, 15.0)159 defaults["gain"] = clip(defaults["gain"] + 0.8, -60.0, 12.0)160 161 if transient > 0.7 and cls in {"fx", "synth_lead"}:162 defaults["az"] += 12.0163 164 if lufs < -24 and cls in {"fx", "ambience", "pad"}:165 defaults["dist"] = clip(defaults["dist"] + 0.5, 0.5, 15.0)166 167 # Add style-specific touch168 sweep = 0.0169 if style == "club" and cls in {"fx", "synth_lead"}:170 sweep = 28.0 if cls == "fx" else 10.0171 elif style in {"cinematic", "film"} and cls in {"fx", "ambience"}:172 sweep = 18.0173 elif style == "podcast" and cls in {"dialogue", "back_vocal"}:174 sweep = 0.0175 176 obj = {177 "id": str(stem.get("id", f"obj_{idx+1}")),178 "class": cls if cls in SPATIAL_CLASSES else "other",179 "az_deg": round(clip(defaults["az"], -180, 180), 2),180 "el_deg": round(clip(defaults["el"], -45, 90), 2),181 "dist_m": round(clip(defaults["dist"], 0.5, 15.0), 2),182 "width": round(clip(defaults["width"], 0.0, 1.0), 2),183 "gain_db": round(clip(defaults["gain"], -60, 12), 2),184 "reverb_send": round(clip(defaults["rev"], 0.0, 1.0), 2),185 "early_reflections": round(clip(defaults["er"], 0.0, 1.0), 2),186 "motion": make_motion(defaults["az"], defaults["el"], defaults["dist"], sweep=sweep),187 }188 189 # Low-end mono safety190 if cls in {"kick", "bass"}:191 obj["az_deg"] = 0.0192 obj["width"] = 0.0 if cls == "kick" else min(obj["width"], 0.08)193 obj["motion"] = make_motion(0.0, obj["el_deg"], obj["dist_m"], sweep=0.0)194 195 objects.append(obj)196 197 scene = {198 "version": "1.0",199 "bed": {200 "layout": target_layout(payload),201 "loudness_target_lufs": -16.0 if payload.get("style") == "cinematic" else -14.0,202 "room_preset": room_preset(payload),203 },204 "objects": objects or [205 {206 "id": "placeholder",207 "class": "other",208 "az_deg": 0.0,209 "el_deg": 0.0,210 "dist_m": 2.0,211 "width": 0.2,212 "gain_db": 0.0,213 "reverb_send": 0.1,214 "early_reflections": 0.1,215 "motion": make_motion(0.0, 0.0, 2.0),216 }217 ],218 "constraints_applied": constraints_applied,219 }220 return scene221 222 223def scene_stats(scene: Dict[str, Any]) -> Dict[str, Any]:224 objects = scene.get("objects", [])225 dominant = sorted(objects, key=lambda o: float(o.get("gain_db", -99.0)), reverse=True)[:3]226 return {227 "layout": scene.get("bed", {}).get("layout", "unknown"),228 "room_preset": scene.get("bed", {}).get("room_preset", "unknown"),229 "object_count": len(objects),230 "dominant": ", ".join(f"{o.get('id')} ({o.get('class')})" for o in dominant) or "n/a",231 }232 233 234def scene_table(scene: Dict[str, Any]) -> List[List[Any]]:235 rows = []236 for obj in scene.get("objects", []):237 rows.append([238 obj.get("id"),239 obj.get("class"),240 obj.get("az_deg"),241 obj.get("el_deg"),242 obj.get("dist_m"),243 obj.get("width"),244 obj.get("gain_db"),245 ])246 return rows247 248 249def scene_markdown(scene: Dict[str, Any], valid: bool, errors: List[str], backend_used: str) -> str:250 stats = scene_stats(scene)251 badge = "✅ Schema valid" if valid else "⚠️ Validation issues"252 lines = [253 f"### {badge}",254 "",255 f"- **Backend:** {backend_used}",256 f"- **Layout:** `{stats['layout']}`",257 f"- **Room preset:** `{stats['room_preset']}`",258 f"- **Objects:** `{stats['object_count']}`",259 f"- **Top objects:** {stats['dominant']}",260 ]261 if errors:262 lines.append("")263 lines.append("**Validation messages**")264 for err in errors[:8]:265 lines.append(f"- {err}")266 return "\n".join(lines)267 268 269def plot_scene(scene: Dict[str, Any]):270 fig = plt.figure(figsize=(7.0, 7.0))271 ax = fig.add_subplot(111)272 ax.set_facecolor("#f8fbff")273 fig.patch.set_facecolor("#f8fbff")274 275 # rings276 for r in [2, 4, 6, 8, 10]:277 circle = plt.Circle((0, 0), r, color="#d9e5f7", fill=False, linewidth=1)278 ax.add_artist(circle)279 280 # axes281 ax.axhline(0, color="#d9e5f7", linewidth=1)282 ax.axvline(0, color="#d9e5f7", linewidth=1)283 284 # labels285 ax.text(0, 10.7, "Front", ha="center", va="bottom", fontsize=11, color="#4a5b77")286 ax.text(10.7, 0, "Right", ha="left", va="center", fontsize=11, color="#4a5b77")287 ax.text(-10.7, 0, "Left", ha="right", va="center", fontsize=11, color="#4a5b77")288 ax.text(0, -10.7, "Rear", ha="center", va="top", fontsize=11, color="#4a5b77")289 290 palette = ["#1d4ed8", "#0891b2", "#7c3aed", "#ea580c", "#0f766e", "#be123c", "#0369a1", "#4d7c0f"]291 292 for idx, obj in enumerate(scene.get("objects", [])):293 az = math.radians(float(obj.get("az_deg", 0.0)))294 dist = float(obj.get("dist_m", 1.0))295 x = dist * math.sin(az)296 y = dist * math.cos(az)297 size = 120 + 220 * float(obj.get("width", 0.2))298 color = palette[idx % len(palette)]299 ax.scatter([x], [y], s=size, c=color, alpha=0.85, edgecolors="white", linewidths=1.5, zorder=3)300 ax.text(x, y + 0.42, f"{obj.get('id')}\n{obj.get('class')}", ha="center", va="bottom", fontsize=9, color="#1f2937")301 302 ax.set_xlim(-11, 11)303 ax.set_ylim(-11, 11)304 ax.set_xticks([])305 ax.set_yticks([])306 ax.set_title("Spatial Scene Preview", fontsize=15, color="#15233d", pad=12)307 return fig308