Spaces:
Running on Zero
Running on Zero
| """Output parsing for model predictions.""" | |
| import json | |
| import re | |
| from typing import Optional | |
| def parse_v5_structured_json( | |
| output: str, | |
| ) -> tuple[list[list[float]], list[Optional[str]]]: | |
| """Parse v5 structured JSON output. | |
| Expected format (inside ```json ... ```): | |
| [{"point_3d": [u, v, d], "label": "1", "gripper": "close"}, ...] | |
| Returns: | |
| (waypoints, actions) where waypoints are [[u, v, depth], ...] and | |
| actions are ["Close Gripper", "Open Gripper", None, ...] | |
| """ | |
| m = re.search(r"```json\s*(.*?)\s*```", output, re.DOTALL) | |
| raw = m.group(1) if m else output.strip() | |
| try: | |
| entries = json.loads(raw) | |
| except json.JSONDecodeError: | |
| return [], [] | |
| waypoints, actions = [], [] | |
| for entry in entries: | |
| if not isinstance(entry, dict): | |
| continue | |
| pt = entry.get("point_3d", entry.get("point_2d")) | |
| if pt is None: | |
| continue | |
| waypoints.append( | |
| [float(pt[0]), float(pt[1]), float(pt[2]) if len(pt) > 2 else 0.0] | |
| ) | |
| grip = entry.get("gripper", "none") | |
| if grip == "close": | |
| actions.append("Close Gripper") | |
| elif grip == "open": | |
| actions.append("Open Gripper") | |
| else: | |
| actions.append(None) | |
| return waypoints, actions | |
| def parse_v3_ans_tags( | |
| output: str, | |
| ) -> tuple[list[list[float]], list[Optional[str]]]: | |
| """Parse v3 format: <ans>[[u, v, d], <action>Close Gripper</action>, ...]</ans>""" | |
| waypoints, actions = [], [] | |
| ans_match = re.search(r"<ans>(.*?)</ans>", output, re.DOTALL) | |
| if not ans_match: | |
| ans_match = re.search(r"\[\[.*?\]\]", output, re.DOTALL) | |
| if not ans_match: | |
| return [], [] | |
| content = ans_match.group(0) | |
| else: | |
| content = ans_match.group(1) | |
| coord_pat = r"\[(\d+(?:\.\d+)?),\s*(\d+(?:\.\d+)?),\s*(\d+(?:\.\d+)?)\]" | |
| action_pat = r"<action>(.*?)</action>" | |
| parts = re.split(action_pat, content) | |
| for i, part in enumerate(parts): | |
| if i % 2 == 0: | |
| for c in re.findall(coord_pat, part): | |
| waypoints.append([float(c[0]), float(c[1]), float(c[2])]) | |
| actions.append(None) | |
| else: | |
| if actions: | |
| actions[-1] = part.strip() | |
| return waypoints, actions | |
| def parse_trajectory( | |
| output: str, prompt_style: str = "v5" | |
| ) -> tuple[list[list[float]], list[Optional[str]]]: | |
| """Auto-detect and parse trajectory from model output. | |
| Args: | |
| output: Raw model output string | |
| prompt_style: "v5" for structured JSON, "v3" for <ans> tags | |
| Returns: | |
| (waypoints, actions) | |
| """ | |
| if prompt_style == "v5": | |
| wp, act = parse_v5_structured_json(output) | |
| if wp: | |
| return wp, act | |
| # Fall back to v3 | |
| return parse_v3_ans_tags(output) | |
| else: | |
| wp, act = parse_v3_ans_tags(output) | |
| if wp: | |
| return wp, act | |
| return parse_v5_structured_json(output) | |