#基于深度学习的滑动验证码缺口识别教程

from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import time
import os

# 自定义配置：收集数量、演示站地址、保存路径
COUNT = 100  # 建议至少200-500张，准度才稳
URL = 'https://captcha1.scrape.center/'
OUTPUT_DIR = 'data/captcha/images'

if not os.path.exists(OUTPUT_DIR):
    os.makedirs(OUTPUT_DIR)

for i in range(COUNT):
    # 每次循环开/关浏览器，避免会话污染导致验证码不刷新
    driver = webdriver.Chrome()
    try:
        driver.get(URL)
        wait = WebDriverWait(driver, 10)
        # 先点登录触发验证码
        login_btn = wait.until(EC.element_to_be_clickable((By.CSS_SELECTOR, '.btn-primary')))
        login_btn.click()
        # 等待背景图完全加载
        captcha_bg = wait.until(EC.presence_of_element_located((By.CSS_SELECTOR, '.geetest_slicebg')))
        time.sleep(2)  # 极验有时候会延迟加载，多等1-2秒
        # 只截图背景图部分，减少无关干扰
        captcha_bg.screenshot(f'{OUTPUT_DIR}/captcha_{i}.png')
    finally:
        driver.quit()

import xmltodict
import os

# 目录配置
XML_DIR = 'data/captcha/annotations'
LABEL_DIR = 'data/captcha/labels'
if not os.path.exists(LABEL_DIR):
    os.makedirs(LABEL_DIR)

def parse_single_xml(xml_path):
    with open(xml_path, 'r', encoding='utf-8') as f:
        xml_dict = xmltodict.parse(f.read())
    
    # 取图片尺寸（用于归一化）
    img_w = int(xml_dict['annotation']['size']['width'])
    img_h = int(xml_dict['annotation']['size']['height'])
    
    # 取缺口的绝对坐标
    bndbox = xml_dict['annotation']['object']['bndbox']
    xmin = int(bndbox['xmin'])
    ymin = int(bndbox['ymin'])
    xmax = int(bndbox['xmax'])
    ymax = int(bndbox['ymax'])
    
    # 转 YOLO 归一化格式
    x_center = ((xmin + xmax) / 2) / img_w
    y_center = ((ymin + ymax) / 2) / img_h
    box_w = (xmax - xmin) / img_w
    box_h = (ymax - ymin) / img_h
    
    return [x_center, y_center, box_w, box_h]

# 批量转换所有XML
for filename in os.listdir(XML_DIR):
    if not filename.endswith('.xml'):
        continue
    # 转换
    yolo_data = parse_single_xml(os.path.join(XML_DIR, filename))
    # 保存为同名txt
    with open(os.path.join(LABEL_DIR, filename.replace('.xml', '.txt')), 'w', encoding='utf-8') as f:
        f.write(f'0 {" ".join(map(str, yolo_data))}')

from fastapi import FastAPI, UploadFile, File, HTTPException
from PIL import Image
import io
import torch
from models.yolo import Model
from utils.general import non_max_suppression, scale_coords
from utils.augmentations import letterbox
import numpy as np

app = FastAPI(title="滑动验证码缺口识别API", version="1.0")

# ------------------ 全局加载模型（只加载一次，避免每次请求都加载） ------------------
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 1. 加载模型结构
model = Model(cfg='models/yolov3.yaml', ch=3, nc=1).to(DEVICE)
# 2. 加载最优权重
checkpoint = torch.load('checkpoints/best.pt', map_location=DEVICE)
model.load_state_dict(checkpoint['model'])
# 3. 切换到评估模式（关闭dropout、batchnorm的训练特性）
model.eval()

def preprocess(image: Image.Image, img_size: int = 640):
    """YOLOv3 标准预处理：letterbox缩放 + 转tensor + 归一化"""
    # 1. 转RGB（防止上传的是灰度图或RGBA图）
    img = image.convert('RGB')
    # 2. letterbox缩放：保持长宽比，填充黑边到指定尺寸
    img = letterbox(img, new_shape=img_size, auto=False)[0]
    # 3. 转 numpy -> (H, W, C) -> (C, H, W) -> 转 torch tensor
    img = np.asarray(img).transpose(2, 0, 1)
    img = torch.from_numpy(img).float().to(DEVICE)
    # 4. 归一化到0-1
    img /= 255.0
    # 5. 加一个batch维度（YOLOv3要求输入是batch的形式）
    if img.ndimension() == 3:
        img = img.unsqueeze(0)
    return img, image.size

@app.post("/predict", summary="识别滑动验证码缺口")
async def predict(file: UploadFile = File(...)):
    # 1. 读取上传的图片
    try:
        image_bytes = await file.read()
        image = Image.open(io.BytesIO(image_bytes))
    except Exception:
        raise HTTPException(status_code=400, detail="上传的不是有效图片")
    
    # 2. 预处理
    img_tensor, original_size = preprocess(image)
    
    # 3. 预测（不计算梯度，加速推理）
    with torch.no_grad():
        pred = model(img_tensor)[0]  # 取第一个输出（检测结果）
        # 4. 非极大值抑制（NMS）：过滤重叠的框，只留置信度最高的
        pred = non_max_suppression(pred, conf_thres=0.5, iou_thres=0.5)[0]
    
    # 5. 后处理：把letterbox缩放后的框坐标还原成原图的
    if pred is not None and len(pred):
        pred = scale_coords(img_tensor.shape[2:], pred[:, :4], original_size).round()
        # 取第一个（也是唯一的）框的坐标
        x1, y1, x2, y2 = pred[0].tolist()
        return {
            "success": True,
            "data": {
                "x1": int(x1),  # 缺口左上角x坐标（滑动距离就是这个值！）
                "y1": int(y1),
                "x2": int(x2),
                "y2": int(y2),
                "width": int(x2 - x1),
                "height": int(y2 - y1)
            }
        }
    else:
        raise HTTPException(status_code=404, detail="未检测到缺口")