x程 点x选x图x标

 

测试 20 张图,  只有一个小图标选择错误, 成功率还行, 就这样吧.

 

"""
某程 图标点选验证码 识别 
pip install -U xmi_logger
"""

import numpy as np
import cv2
from PIL import Image, ImageDraw
from pathlib import Path
from typing import List, Tuple, Dict, Optional
import io
import base64
import itertools

from xmi_logger import XmiLogger
logger = XmiLogger("xc_0")


GRID_X = [42, 85, 127, 170, 213, 255]
GRID_Y_RANGE = (35, 145)
ICON_HALF = 24             



def _load_image(src) -> Image.Image:
    if isinstance(src, Image.Image):
        return src
    if isinstance(src, (str, Path)):
        p = str(src)
        if p.startswith("data:"):
            return Image.open(io.BytesIO(base64.b64decode(p.split(",", 1)[1])))
        if len(p) > 260:
            try:
                return Image.open(io.BytesIO(base64.b64decode(p)))
            except Exception:
                pass
        return Image.open(p)
    if isinstance(src, bytes):
        return Image.open(io.BytesIO(src))
    raise ValueError(f"Unsupported: {type(src)}")



def _get_alpha_map(big_bgr: np.ndarray) -> np.ndarray:
    big_bgr_f = big_bgr.astype(np.float32)
    bg_estimated = cv2.GaussianBlur(big_bgr_f, (51, 51), 0)
    C_overlay = np.array([120, 107, 93], dtype=np.float32) # BGR
    
    dist_out = np.sqrt(np.sum((big_bgr_f - C_overlay) ** 2, axis=2))
    dist_bg = np.sqrt(np.sum((bg_estimated - C_overlay) ** 2, axis=2))
    
    dist_bg_safe = np.where(dist_bg < 1.0, 1.0, dist_bg)
    ratio = dist_out / dist_bg_safe
    
    alpha = np.clip(1.0 - ratio, 0, 1)
    return (alpha * 255).astype(np.uint8)


def _get_clean_templates(small_pil: Image.Image, n_icons: int = 3) -> List[np.ndarray]:
    sarr = np.array(small_pil.convert("RGBA"))
    w = small_pil.size[0]
    cell_w = w / n_icons

    templates = []
    for i in range(n_icons):
        left = int(round(i * cell_w))
        right = int(round((i + 1) * cell_w))
        alpha = sarr[:, left:right, 3]

        coords = np.argwhere(alpha > 50)
        if len(coords) == 0:
            templates.append(np.zeros((1, 1), dtype=np.uint8))
            continue
        y0, x0 = coords.min(axis=0)
        y1, x1 = coords.max(axis=0)
        templates.append(alpha[y0:y1 + 1, x0:x1 + 1])

    return templates


def _get_crop_32x32(img: np.ndarray) -> np.ndarray:
    if img is None or img.size == 0:
        return np.zeros((32, 32), dtype=np.uint8)
    _, bin_img = cv2.threshold(img, 50, 255, cv2.THRESH_BINARY)
    if bin_img is None:
        return np.zeros((32, 32), dtype=np.uint8)
    coords = np.argwhere(bin_img > 0)
    if len(coords) == 0:
        return np.zeros((32, 32), dtype=np.uint8)
    y0, x0 = coords.min(axis=0)
    y1, x1 = coords.max(axis=0)
    cropped = img[y0:y1+1, x0:x1+1]
    return cv2.resize(cropped, (32, 32), interpolation=cv2.INTER_AREA)


def _compute_global_similarity(c_img: np.ndarray, t_img: np.ndarray) -> float:
    c_f = c_img.astype(np.float32)
    t_f = t_img.astype(np.float32)
    
    res_pos = cv2.matchTemplate(c_f, t_f, cv2.TM_CCOEFF_NORMED)[0, 0]
    res_neg = cv2.matchTemplate(255.0 - c_f, 255.0 - t_f, cv2.TM_CCOEFF_NORMED)[0, 0]
    
    return float(res_pos * 0.5 + res_neg * 0.5)


@logger.time_it
def solve_xiecheng(
    bg_input,
    small_input,
    *,
    return_annotated: bool = True,
    n_icons: int = 3,
) -> Dict:
    bg_pil = _load_image(bg_input).convert("RGB")
    small_pil = _load_image(small_input).convert("RGBA")

    bg_arr = np.array(bg_pil)
    bg_bgr = cv2.cvtColor(bg_arr, cv2.COLOR_RGB2BGR)

    alpha_map = _get_alpha_map(bg_bgr)

    templates = _get_clean_templates(small_pil, n_icons)
    templates_32 = [_get_crop_32x32(t) for t in templates]

    score_matrix = {}
    for ii, template in enumerate(templates):
        th, tw = template.shape
        t_aspect = tw / th
        t_img_32 = templates_32[ii]

        for col_idx, gx in enumerate(GRID_X):
            x_lo = max(0, gx - 15)
            x_hi = min(alpha_map.shape[1], gx + 15)
            y_lo = max(0, GRID_Y_RANGE[0] - 10)
            y_hi = min(alpha_map.shape[0], GRID_Y_RANGE[1] + 10)

            region = alpha_map[y_lo:y_hi, x_lo:x_hi]

            best_match_score = -1.0
            best_loc = (0, 0)
            best_nw, best_nh = 0, 0

            for nw in range(24, 46, 2):
                for nh in range(20, 42, 2):
                    resized_t = cv2.resize(template, (nw, nh), interpolation=cv2.INTER_AREA)
                    if region.shape[0] < nh or region.shape[1] < nw:
                        continue

                    res = cv2.matchTemplate(region, resized_t, cv2.TM_CCOEFF_NORMED)
                    _, max_val, _, max_loc = cv2.minMaxLoc(res)

                    if max_val > best_match_score:
                        best_match_score = max_val
                        best_loc = max_loc
                        best_nw = nw
                        best_nh = nh

            tx, ty = best_loc
            cand_crop = region[ty:ty+best_nh, tx:tx+best_nw]
            cand_crop_32 = _get_crop_32x32(cand_crop)
            glob_score = _compute_global_similarity(cand_crop_32, t_img_32)

            cand_aspect = best_nw / best_nh
            if abs(cand_aspect - t_aspect) > 0.45:
                final_score = -1.0
            else:
                final_score = glob_score

            cx = x_lo + tx + best_nw // 2
            cy = y_lo + ty + best_nh // 2
            score_matrix[(ii, col_idx)] = (final_score, cx, cy, best_nw)

    best_perm = None
    best_total_score = -1.0
    for perm in itertools.permutations(range(6), n_icons):
        score_sum = sum(score_matrix[(i, col_idx)][0] for i, col_idx in enumerate(perm))
        if score_sum > best_total_score:
            best_total_score = score_sum
            best_perm = perm

    coordinates = []
    details = []
    for i, col_idx in enumerate(best_perm):
        score, cx, cy, cell = score_matrix[(i, col_idx)]
        coordinates.append((cx, cy))
        details.append({
            "center": (cx, cy),
            "score": score,
            "cell": cell,
            "bbox": (cx - cell // 2, cy - cell // 2, cell, cell),
            "method": "alpha_global_shape_matching"
        })

    result: Dict = {
        "coordinates": coordinates,
        "details": details,
        "annotated_image": None,
    }

    if return_annotated:
        result["annotated_image"] = _draw_annotation(
            bg_pil, small_pil, coordinates, details, n_icons
        )
    logger.info(f"Result: {result}")
    return result



def _draw_annotation(
    bg_pil: Image.Image,
    small_pil: Image.Image,
    coordinates: List,
    details: List,
    n_icons: int = 3,
) -> Image.Image:
    colors = [
        (255, 50, 50), (50, 220, 50), (50, 100, 255),
        (255, 255, 0), (255, 0, 255), (0, 255, 255),
    ]

    small_rgb = small_pil.convert("RGBA")
    sw, sh = small_rgb.size
    scale_factor = 3 
    small_resized = small_rgb.resize((sw * scale_factor, sh * scale_factor),
                                     Image.NEAREST)
    sr_w, sr_h = small_resized.size

    small_canvas = Image.new("RGB", (sr_w, sr_h + 16), (30, 30, 30))
    small_canvas.paste(small_resized, (0, 16), small_resized)
    small_draw = ImageDraw.Draw(small_canvas)

    cell_w = sw * scale_factor / n_icons
    for i in range(n_icons):
        color = colors[i % len(colors)]
        left = int(round(i * cell_w))
        right = int(round((i + 1) * cell_w)) - 1
        small_draw.rectangle([left, 16, right, sr_h + 15], outline=color, width=2)
        label = f"icon{i}"
        small_draw.text((left + 4, 2), label, fill=color)
        if i > 0:
            x = int(round(i * cell_w))
            small_draw.line([(x, 16), (x, sr_h + 15)], fill=(200, 200, 200), width=1)

    annotated = bg_pil.copy()
    draw = ImageDraw.Draw(annotated)

    for i, (coord, detail) in enumerate(zip(coordinates, details)):
        if not coord or not detail:
            continue
        cx, cy = coord
        color = colors[i % len(colors)]

        if "bbox" in detail:
            bx, by, bw, bh = detail["bbox"]
            draw.rectangle([bx, by, bx + bw, by + bh], outline=color, width=2)
        else:
            draw.rectangle([cx - 20, cy - 20, cx + 20, cy + 20], outline=color, width=2)

        draw.line([cx - 6, cy, cx + 6, cy], fill=color, width=2)
        draw.line([cx, cy - 6, cx, cy + 6], fill=color, width=2)

        label = f"icon{i} ({detail['score']:.2f})"
        label_y = detail.get("bbox", (0, cy - 20, 0, 0))[1]
        draw.text((cx - 15, max(0, label_y - 14)), label, fill=color)

    bg_w, bg_h = annotated.size
    sc_w, sc_h = small_canvas.size
    total_w = max(bg_w, sc_w)
    total_h = bg_h + sc_h + 4 
    combined = Image.new("RGB", (total_w, total_h), (30, 30, 30))
    combined.paste(annotated, ((total_w - bg_w) // 2, 0))
    combined.paste(small_canvas, ((total_w - sc_w) // 2, bg_h + 4))

    return combined



def solve_xiecheng_api(bg_input, small_input, **kwargs) -> Dict:
    result = solve_xiecheng(bg_input, small_input, return_annotated=False, **kwargs)
    return {"coordinates": result["coordinates"], "details": result["details"]}



def process_all():
    script_dir = Path(__file__).resolve().parent
    big_dir = script_dir / "images_raw"
    small_dir = script_dir / "small_img"
    output_dir = script_dir / "debug"
    output_dir.mkdir(exist_ok=True)


    big_files = sorted(big_dir.glob("big_image_*.png"))
    logger.info(f"Found {len(big_files)} big images\n")

    total = 0
    correct_count = 0

    for bf in big_files:
        ts = bf.stem.replace("big_image_", "")
        sf = small_dir / f"small_image_{ts}.png"
        if not sf.exists():
            logger.info(f"[{ts}] Skip: small image not found")
            continue

        total += 1
        result = solve_xiecheng(str(bf), str(sf), return_annotated=True)

        pred_cols = []
        for c in result["coordinates"]:
            cx = c[0]
            col_idx = min(range(len(GRID_X)), key=lambda idx: abs(GRID_X[idx] - cx))
            pred_cols.append(col_idx)
        
        pred_perm = tuple(pred_cols)
        
        for i, (c, d) in enumerate(zip(result["coordinates"], result["details"])):
            logger.info(f"  icon{i}: center={c}  score={d['score']:.3f}  method={d['method']}")

        if result["annotated_image"]:
            result["annotated_image"].save(output_dir / f"{ts}_annotated.png")

    if total:
        accuracy = (correct_count / total) * 100
        logger.info(f"Total Evaluated: {total}  Correct: {correct_count}  "
              f"Accuracy Rate: {accuracy:.1f}%")
    logger.info(f"Results saved to {output_dir}/")


if __name__ == "__main__":
    process_all()