예제 소스

probe_move.py

from __future__ import annotations

import argparse
import io
import os
import socket
import struct
from dataclasses import dataclass
from typing import Optional, Tuple

from PIL import ImageDraw

from .dataset_baduk import BOARD_SIZE, _GoSim, render_board_png
from .sgf_parser import parse_sgf_file

MAGIC_GO19 = b"GO19"
MAGIC_RPLY = b"RPLY"


def _u32be(v: int) -> bytes:
    return struct.pack(">I", v & 0xFFFFFFFF)


def _recv_exact(sock: socket.socket, n: int) -> bytes:
    buf = bytearray()
    while len(buf) < n:
        chunk = sock.recv(n - len(buf))
        if not chunk:
            raise EOFError("peer closed")
        buf.extend(chunk)
    return bytes(buf)


def _send_go19(sock: socket.socket, png_bytes: bytes, board_bytes: bytes, ai_side: int) -> None:
    if len(board_bytes) != BOARD_SIZE * BOARD_SIZE:
        raise ValueError(f"board_bytes must be {BOARD_SIZE*BOARD_SIZE} bytes")
    if ai_side not in (1, 2):
        raise ValueError("ai_side must be 1(B) or 2(W)")
    body = bytearray()
    body.append(BOARD_SIZE)
    body.append(ai_side)
    body += _u32be(len(png_bytes))
    body += png_bytes
    body += board_bytes
    sock.sendall(MAGIC_GO19 + _u32be(len(body)) + body)


def _recv_reply(sock: socket.socket) -> Tuple[bool, Optional[int], Optional[int], Optional[str]]:
    head = _recv_exact(sock, 4)
    if head != MAGIC_RPLY:
        raise ValueError(f"bad reply header: {head!r}")
    flag = _recv_exact(sock, 1)[0]
    if flag == 1:
        row = _recv_exact(sock, 1)[0]
        col = _recv_exact(sock, 1)[0]
        if row == 0xFF and col == 0xFF:
            return True, None, None, None
        return True, int(row), int(col), None
    if flag == 0:
        (n,) = struct.unpack(">H", _recv_exact(sock, 2))
        msg = _recv_exact(sock, n).decode("utf-8", errors="replace")
        return False, None, None, msg
    raise ValueError(f"unknown flag: {flag}")


def _count_stones(board_bytes: bytes) -> int:
    # 0=empty, 1=black, 2=white
    return sum(1 for b in board_bytes if b)


def _build_position_from_sgf(sgf_path: str, ply: int) -> Tuple[bytes, int, int, int, int]:
    """Return (board_bytes, next_side_int, moves_total, moves_applied, stones_count)."""
    g = parse_sgf_file(sgf_path)
    if g is None:
        raise SystemExit(f"SGF 파싱 실패: {sgf_path}")
    if g.size != BOARD_SIZE:
        raise SystemExit(f"현재 probe는 19x19만 지원합니다 (SGF size={g.size})")

    sim = _GoSim(size=BOARD_SIZE)
    for (r, c) in g.setup_black:
        sim.set(r, c, 1)
    for (r, c) in g.setup_white:
        sim.set(r, c, 2)

    # apply moves up to ply
    # ply == -1: 마지막 수까지 적용(최종 국면)
    ply_i = int(ply)
    if ply_i < 0:
        n = len(g.moves)
    else:
        n = max(0, min(ply_i, len(g.moves)))
    applied = 0
    for i in range(n):
        color, pt = g.moves[i]
        who = 1 if color.upper() == "B" else 2
        if pt is None:
            continue
        rr, cc = pt
        if sim.play(rr, cc, who):
            applied += 1

    next_side = 1
    if n > 0:
        last_color = g.moves[n - 1][0].upper()
        next_side = 2 if last_color == "B" else 1
    board = sim.packed()
    stones = _count_stones(board)
    return board, next_side, len(g.moves), applied, stones


def main(argv: Optional[list[str]] = None) -> int:
    p = argparse.ArgumentParser(prog="python -m vlm_server probe", description="GO19 서버에 한 수 질의")
    p.add_argument("--host", default="127.0.0.1")
    p.add_argument("--port", type=int, default=9319)
    p.add_argument("--sgf", required=True, help="국면을 가져올 SGF 파일 경로")
    p.add_argument("--ply", type=int, default=-1, help="-1=최종, 0=초기, 1=첫 수 후, ...")
    p.add_argument(
        "--side",
        choices=["auto", "B", "W"],
        default="auto",
        help="질의할 차례(흑/백). auto면 sgf ply로 다음 차례 추정",
    )
    p.add_argument("--out-png", default=None, help="렌더된 보드 PNG 저장 경로(선택)")
    args = p.parse_args(argv)

    board_bytes, next_side, moves_total, moves_applied, stones = _build_position_from_sgf(
        args.sgf, args.ply
    )
    ai_side = next_side
    if args.side == "B":
        ai_side = 1
    elif args.side == "W":
        ai_side = 2

    if args.ply < 0 and stones == 0 and moves_total > 0:
        print(
            f"WARN: ply=-1 인데 돌이 0개입니다. SGF 적용이 거의 실패했을 수 있습니다. "
            f"(moves_total={moves_total}, moves_applied={moves_applied})"
        )

    img = render_board_png(board_bytes, BOARD_SIZE, px=768)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    png_bytes = buf.getvalue()

    endpoint = (args.host, int(args.port))
    try:
        sock_ctx = socket.create_connection(endpoint, timeout=10.0)
    except ConnectionRefusedError:
        print(
            f"ERR: 연결 거부 {endpoint[0]}:{endpoint[1]} — GO19 서버가 안 떠 있습니다.\n"
            f"     다른 터미널에서 예: python -m vlm_server serve --port {endpoint[1]}"
        )
        return 3
    except OSError as e:
        print(f"ERR: 소켓 연결 실패 {endpoint}: {e}")
        return 3

    with sock_ctx as sock:
        sock.settimeout(120.0)
        _send_go19(sock, png_bytes, board_bytes, ai_side=ai_side)
        ok, row, col, err = _recv_reply(sock)
        if not ok:
            print(f"ERR: {err}")
            return 2
        move_text = "PASS" if (row is None or col is None) else f"row={row} col={col}"

        # out png에 추천 수 마킹까지 해서 저장
        if args.out_png:
            out_img = img.copy()
            d = ImageDraw.Draw(out_img)
            if row is not None and col is not None:
                px = out_img.size[0]
                pad = px * 0.045
                avail = px - 2.0 * pad
                sp = avail / (BOARD_SIZE - 1)
                cx = pad + col * sp
                cy = pad + row * sp
                r = max(6.0, sp * 0.22)
                d.ellipse([(cx - r, cy - r), (cx + r, cy + r)], outline=(230, 40, 40), width=5)
            # 좌상단에 텍스트로도 남김 (파일만 봐도 어떤 테스트인지 알 수 있게)
            side_name = "B" if ai_side == 1 else "W"
            label = f"{os.path.basename(args.sgf)} ply={args.ply} side={side_name} {move_text} stones={stones}"
            d.rectangle([(6, 6), (out_img.size[0] - 6, 34)], fill=(255, 255, 255, 200))
            d.text((10, 10), label, fill=(10, 10, 10))
            out_img.save(args.out_png, format="PNG")

        print(
            f"OK: {move_text} (stones={stones}, moves_applied={moves_applied}/{moves_total}, ply={args.ply})"
        )
        return 0


if __name__ == "__main__":
    raise SystemExit(main())