X_SSL_Net/lib/data/ddti.py

from __future__ import annotations

import json
from pathlib import Path
import xml.etree.ElementTree as ET

from PIL import Image, ImageDraw


def parse_ddti_xml(annotation_path: str | Path) -> dict[int, list[list[tuple[int, int]]]]:
    """
    解析 DDTI 的 xml 标注。

    Returns:
        {image_index: [polygon1, polygon2, ...]}
    """
    annotation_path = Path(annotation_path)
    root = ET.parse(annotation_path).getroot()
    image_to_polygons: dict[int, list[list[tuple[int, int]]]] = {}

    for mark in root.findall("mark"):
        image_text = mark.findtext("image")
        svg_text = mark.findtext("svg")
        if not image_text or not svg_text:
            continue

        image_index = int(image_text)
        try:
            shapes = json.loads(svg_text)
        except json.JSONDecodeError:
            continue

        polygons: list[list[tuple[int, int]]] = []
        for shape in shapes:
            points = shape.get("points", [])
            polygon = []
            for point in points:
                x = int(round(point["x"]))
                y = int(round(point["y"]))
                polygon.append((x, y))
            if len(polygon) >= 3:
                polygons.append(polygon)

        if polygons:
            image_to_polygons[image_index] = polygons

    return image_to_polygons


def build_ddti_mask(
        image_path: str | Path,
        annotation_path: str | Path,
        image_index: int | None = None,
        fill_value: int = 255,
) -> Image.Image:
    """
    根据 DDTI 的 xml 为指定图像生成二值掩膜。
    """
    image_path = Path(image_path)
    annotation_path = Path(annotation_path)
    image = Image.open(image_path)
    width, height = image.size

    if image_index is None:
        stem = image_path.stem
        if "_" not in stem:
            raise ValueError(f"Cannot infer image index from file name: {image_path.name}")
        _, image_idx_str = stem.split("_", 1)
        image_index = int(image_idx_str)

    polygons_map = parse_ddti_xml(annotation_path)
    polygons = polygons_map.get(int(image_index), [])

    mask = Image.new("L", (width, height), 0)
    draw = ImageDraw.Draw(mask)
    for polygon in polygons:
        draw.polygon(polygon, outline=fill_value, fill=fill_value)
    return mask


__all__ = ["parse_ddti_xml", "build_ddti_mask"]