#!/usr/bin/env python3
"""标注预览图生成: 动物红色实线框 + 生境 cover 黄色虚线框

用法:
  venv/bin/python visualize_labels.py --input datasets/images --labels datasets/labels --output datasets/pheasant_label
"""
import argparse
from pathlib import Path

from PIL import Image, ImageDraw

CLASSES = ["pheasant", "cover"]
LABELS = {"pheasant": "野鸡", "cover": "疑似区域"}


def draw_dashed(draw, box, outline, width, dash=12, gap=8):
    x1, y1, x2, y2 = box
    for (ax, ay, bx, by) in [(x1, y1, x2, y1), (x2, y1, x2, y2),
                             (x2, y2, x1, y2), (x1, y2, x1, y1)]:
        length = max(abs(bx - ax), abs(by - ay))
        steps = max(int(length / (dash + gap)), 1)
        for i in range(steps):
            s = i / steps
            e = min((i * (dash + gap) + dash) / length, 1.0)
            if e <= s:
                continue
            draw.line([ax + (bx - ax) * s, ay + (by - ay) * s,
                       ax + (bx - ax) * e, ay + (by - ay) * e],
                      fill=outline, width=width)


def main():
    ap = argparse.ArgumentParser(description="标注预览图生成")
    ap.add_argument("--input", default="datasets/images")
    ap.add_argument("--labels", default="datasets/labels")
    ap.add_argument("--output", default="datasets/pheasant_label")
    args = ap.parse_args()

    img_dir = Path(args.input)
    lab_dir = Path(args.labels)
    out_dir = Path(args.output)
    out_dir.mkdir(parents=True, exist_ok=True)

    count = 0
    for img_path in sorted(img_dir.rglob("*")):
        if img_path.suffix.lower() not in (".jpg", ".jpeg", ".png", ".webp", ".bmp"):
            continue
        rel = img_path.relative_to(img_dir)
        lab_path = lab_dir / rel.with_suffix(".txt")
        if not lab_path.exists():
            continue
        im = Image.open(img_path).convert("RGB")
        W, H = im.size
        d = ImageDraw.Draw(im)
        for line in lab_path.read_text().splitlines():
            parts = line.split()
            if not parts:
                continue
            cid, cx, cy, w, h = map(float, parts)
            box = ((cx - w / 2) * W, (cy - h / 2) * H,
                   (cx + w / 2) * W, (cy + h / 2) * H)
            if int(cid) == 1:  # cover 生境: 黄色虚线
                draw_dashed(d, box, (255, 200, 0), width=5)
                d.text((box[0] + 6, max(box[1] - 28, 4)), "疑似区域",
                       fill=(255, 200, 0))
            else:              # 动物: 红色实线
                d.rectangle(box, outline=(255, 0, 0), width=6)
                label = LABELS.get(CLASSES[int(cid)], CLASSES[int(cid)])
                d.text((box[0] + 6, max(box[1] - 28, 4)), label,
                       fill=(255, 0, 0))
        out = out_dir / (img_path.stem + ".jpg")
        im.save(out, quality=92)
        count += 1
        print(out.name)
    print(f"共生成 {count} 张")


if __name__ == "__main__":
    main()