| |
| from __future__ import annotations |
|
|
| import argparse |
| import re |
| from pathlib import Path |
|
|
| from PIL import Image, ImageDraw, ImageFont |
|
|
| from puker_judge_utils import ( |
| BASE_MODEL_ID, |
| MODEL_REPO_ID, |
| encode_image_prompt, |
| generate_answer, |
| load_adapter, |
| print_json, |
| rank_prompt, |
| ) |
|
|
| BOARD_SIZE = (1280, 820) |
| CANDIDATE_SIZE = (600, 360) |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser( |
| description="Choose the best of two to four playing-card candidates." |
| ) |
| parser.add_argument( |
| "candidates", |
| type=Path, |
| nargs="*", |
| help="Two to four rectified candidate images.", |
| ) |
| parser.add_argument( |
| "--board-image", |
| type=Path, |
| help="Use an already constructed board instead of candidate files.", |
| ) |
| parser.add_argument( |
| "--candidate-count", |
| type=int, |
| help="Required with --board-image; must be between 2 and 4.", |
| ) |
| parser.add_argument( |
| "--board-output", |
| type=Path, |
| default=Path("candidate_board.jpg"), |
| ) |
| parser.add_argument("--repo-id", default=MODEL_REPO_ID) |
| parser.add_argument("--base-model-id", default=BASE_MODEL_ID) |
| parser.add_argument( |
| "--int4", |
| action="store_true", |
| help="Use bitsandbytes NF4 weights with BF16 compute.", |
| ) |
| return parser.parse_args() |
|
|
|
|
| def load_font(size: int) -> ImageFont.ImageFont: |
| for path in ( |
| "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", |
| "DejaVuSans-Bold.ttf", |
| ): |
| try: |
| return ImageFont.truetype(path, size=size) |
| except OSError: |
| continue |
| return ImageFont.load_default() |
|
|
|
|
| def build_board(candidate_paths: list[Path]) -> Image.Image: |
| count = len(candidate_paths) |
| if not 2 <= count <= 4: |
| raise ValueError("Provide exactly 2, 3, or 4 candidate images.") |
|
|
| board = Image.new("RGB", BOARD_SIZE, (226, 231, 237)) |
| draw = ImageDraw.Draw(board) |
| font = load_font(34) |
| outer_margin = 28 |
| column_gap = 24 |
| row_gap = 24 |
| cell_width = (BOARD_SIZE[0] - 2 * outer_margin - column_gap) // 2 |
| cell_height = (BOARD_SIZE[1] - 2 * outer_margin - row_gap) // 2 |
| label_height = 46 |
|
|
| for index, path in enumerate(candidate_paths): |
| column = index % 2 |
| row = index // 2 |
| x0 = outer_margin + column * (cell_width + column_gap) |
| y0 = outer_margin + row * (cell_height + row_gap) |
| x1 = x0 + cell_width |
| y1 = y0 + cell_height |
| draw.rounded_rectangle( |
| (x0, y0, x1, y1), |
| radius=14, |
| fill=(244, 246, 248), |
| outline=(178, 184, 192), |
| width=2, |
| ) |
| draw.rounded_rectangle( |
| (x0 + 12, y0 + 8, x0 + 76, y0 + label_height), |
| radius=10, |
| fill=(255, 218, 72), |
| outline=(45, 48, 52), |
| width=2, |
| ) |
| draw.text( |
| (x0 + 44, y0 + 8 + label_height // 2), |
| str(index + 1), |
| font=font, |
| fill=(25, 31, 42), |
| anchor="mm", |
| ) |
| with Image.open(path) as source: |
| candidate = source.convert("RGB").resize( |
| CANDIDATE_SIZE, |
| Image.Resampling.LANCZOS, |
| ) |
| available_width = cell_width - 36 |
| available_height = cell_height - label_height - 28 |
| candidate.thumbnail( |
| (available_width, available_height), |
| Image.Resampling.LANCZOS, |
| ) |
| paste_x = x0 + (cell_width - candidate.width) // 2 |
| paste_y = y0 + label_height + ( |
| cell_height - label_height - candidate.height |
| ) // 2 |
| draw.rectangle( |
| ( |
| paste_x + 4, |
| paste_y + 5, |
| paste_x + 4 + candidate.width, |
| paste_y + 5 + candidate.height, |
| ), |
| fill=(185, 190, 196), |
| ) |
| board.paste(candidate, (paste_x, paste_y)) |
| return board |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| if args.board_image is not None: |
| if args.candidates: |
| raise SystemExit( |
| "Use either candidate files or --board-image, not both." |
| ) |
| if args.candidate_count is None or not 2 <= args.candidate_count <= 4: |
| raise SystemExit( |
| "--candidate-count 2..4 is required with --board-image." |
| ) |
| with Image.open(args.board_image) as source: |
| board = source.convert("RGB").copy() |
| candidate_count = args.candidate_count |
| candidate_paths: list[Path] = [] |
| else: |
| if args.candidate_count is not None: |
| raise SystemExit( |
| "--candidate-count is inferred when candidate files are used." |
| ) |
| candidate_paths = args.candidates |
| candidate_count = len(candidate_paths) |
| board = build_board(candidate_paths) |
| args.board_output.parent.mkdir(parents=True, exist_ok=True) |
| board.save(args.board_output, "JPEG", quality=96, subsampling=0) |
|
|
| model, processor, device, timings = load_adapter( |
| "rank_adapter", |
| repo_id=args.repo_id, |
| base_model_id=args.base_model_id, |
| int4=args.int4, |
| ) |
| inputs = encode_image_prompt( |
| processor, |
| board, |
| rank_prompt(candidate_count), |
| device, |
| ) |
| raw_output, generation_seconds = generate_answer( |
| model, |
| processor, |
| inputs, |
| ) |
| match = re.search(r"[1-4]", raw_output) |
| if match is None or int(match.group(0)) > candidate_count: |
| raise SystemExit(f"Model returned an invalid answer: {raw_output!r}") |
| label = int(match.group(0)) |
| payload = { |
| "selected_label": label, |
| "raw_output": raw_output, |
| "candidate_count": candidate_count, |
| "quantization": "int4-nf4" if args.int4 else "bf16", |
| "generation_seconds": round(generation_seconds, 4), |
| **{key: round(value, 4) for key, value in timings.items()}, |
| } |
| if candidate_paths: |
| payload["selected_file"] = str(candidate_paths[label - 1].resolve()) |
| payload["board_image"] = str(args.board_output.resolve()) |
| else: |
| payload["board_image"] = str(args.board_image.resolve()) |
| print_json(payload) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
