| """Parsing and visualisation of QuPAINT responses. |
| |
| QuPAINT answers in free text and closes with a `<CONCLUSION>` span that holds the |
| flakes it commits to, each as a `<box>x, y, w, h</box>` quadruple in **percent** |
| of image width/height (top-left origin). |
| """ |
|
|
| import re |
| from typing import List, Sequence, Tuple |
|
|
| CONCLUSION_RE = re.compile(r"<\s*/?\s*CONCLUSION>(.*?)<\s*/?\s*CONCLUSION>", re.DOTALL) |
| BOX_RE = re.compile(r"<box>(.*?)</box>", re.DOTALL) |
|
|
| Box = Tuple[float, float, float, float] |
|
|
|
|
| def parse_conclusion_boxes(response: str) -> List[Box]: |
| """Return the `<box>` quadruples inside the `<CONCLUSION>` span, in percent.""" |
| match = CONCLUSION_RE.search(response) |
| if not match: |
| return [] |
| boxes = [] |
| for raw in BOX_RE.findall(match.group(1)): |
| try: |
| coords = [float(v.strip()) for v in raw.split(",")] |
| except ValueError: |
| continue |
| if len(coords) == 4: |
| boxes.append(tuple(coords)) |
| return boxes |
|
|
|
|
| def denorm_box(box: Box, height: int, width: int) -> Tuple[int, int, int, int]: |
| """Percent box (x, y, w, h) -> pixel box (x, y, w, h) for a `height` x `width` image.""" |
| x, y, w, h = box |
| return ( |
| int(x / 100 * width), |
| int(y / 100 * height), |
| int(w / 100 * width), |
| int(h / 100 * height), |
| ) |
|
|
|
|
| def parse_boxes(response: str, height: int, width: int) -> List[Tuple[int, int, int, int]]: |
| """Parse a response straight into pixel boxes for an image of the given size.""" |
| return [denorm_box(box, height, width) for box in parse_conclusion_boxes(response)] |
|
|
|
|
| def draw_boxes( |
| image, |
| boxes: Sequence[Tuple[int, int, int, int]], |
| color: Tuple[int, int, int] = (255, 255, 0), |
| thickness: int = 2, |
| ): |
| """Draw pixel boxes onto a PIL image and return a new annotated image.""" |
| from PIL import ImageDraw |
|
|
| annotated = image.convert("RGB").copy() |
| draw = ImageDraw.Draw(annotated) |
| for x, y, w, h in boxes: |
| draw.rectangle([x, y, x + w, y + h], outline=color, width=thickness) |
| return annotated |
|
|
