rotate-captcha-crack项目重新训练百度旋转验证码角度预测模型

参考：

building-powerful-image-classification-models-using-very-little-data.html

https://github.com/Starry-OvO/rotate-captcha-crack （主）作者思路：https://www.52pojie.cn/thread-1754224-1-1.html

纠正 新版百度、百家号旋转验证码识别

d4net作者博客

训练

图片来源

角度为0的百度验证码图片
可以先爬虫获取多张，然后计算相似度删除重复图片

训练数据集规模

一张验证码可以复制多次（具体多少 看你底图数据量），这样一个epoch内会出现多个角度，模型快速学习，方便收敛

一张验证码不进行多次复制，验证集loss震荡会很厉害，300张图片loss可以降低到1.4 ，但是泛化性也很差

加载原始模型代码

（rotate-captcha-crack原始加载模型方法可能报错，所以重写）

def fineturn_from_old_model(self):

        model_path = r'models/RotNetR/230308_08_02_34_000/best.pth'
        print(model_path,'------------------------------------------')
        state_dict = torch.load(model_path, map_location=torch.device('cuda'))
        model = RotNetR(180)
        model = model.to(device)
        model.load_state_dict(state_dict)
        # 冻结除最后一层之外的所有层
        for name, param in model.named_parameters():
            if not name.startswith('backbone.fc'):  # fc为最后一层的名称
                param.requires_grad = False
            else:
                print(name)
        self.model.load_state_dict(model.state_dict())

数据增强办法

rotate_captcha_crack\dataset\rot.py 中 内容如下：
使用的数据增强办法 ：

• 改变图像颜色的四个方面：亮度、对比度、饱和度和色调
• 随机剪裁
• 左右翻转

from typing import Tuple

import torch
from torch import Tensor
from torch.utils.data import Dataset
from torchvision.transforms import Normalize
from torchvision import transforms

from ..const import DEFAULT_CLS_NUM, DEFAULT_TARGET_SIZE
from .helper import DEFAULT_NORM, from_img
from .typing import TypeImgTsSeq

TypeRotItem = Tuple[Tensor, Tensor]


class RotDataset(Dataset[TypeRotItem]):
    """
    Dataset for RotNet (classification).

    Args:
        imgseq (TypeImgSeq): upstream dataset
        target_size (int, optional): output img size. Defaults to `DEFAULT_TARGET_SIZE`.
        norm (Normalize, optional): normalize policy. Defaults to `DEFAULT_NORM`.

    Methods:
        - `def __len__(self) -> int:` length of the dataset
        - `def __getitem__(self, idx: int) -> TypeRotItem:` get square img_ts and index_ts
            ([C,H,W]=[3,target_size,target_size], dtype=float32, range=[0.0,1.0)), ([N]=[1], dtype=long, range=[0,cls_num))
    """

    __slots__ = [
        'imgseq',
        'cls_num',
        'target_size',
        'norm',
        'size',
        'indices',
    ]

    def __init__(
        self,
        imgseq: TypeImgTsSeq,
        cls_num: int = DEFAULT_CLS_NUM,
        target_size: int = DEFAULT_TARGET_SIZE,
        norm: Normalize = DEFAULT_NORM,
    ) -> None:
        self.imgseq = imgseq
        self.cls_num = cls_num
        self.target_size = target_size
        self.norm = norm

        self.size = self.imgseq.__len__()
        self.indices = torch.randint(cls_num, (self.size,), dtype=torch.long)
        self.transforms = transforms.Compose([
                    # transforms.Resize(240), # 将图像最短边缩至240，宽高比例不变
                    transforms.RandomHorizontalFlip(), # 以0.5的概率左右翻转图像
                    # transforms.ToTensor(), # 将PIL图像转为Tensor，并且进行归一化
                    # transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) # 进行mean与std为0.5的标准化
                     transforms.RandomResizedCrop( (350, 350), scale=(0.8, 1), ratio=(0.5, 2)),
               # 随机裁剪一个面积为原始面积50%到100%的区域，该区域的宽高比从0.5～2之间随机取值。 然后，区域的宽度和高度都被缩放到350像素。
                     # 我们可以改变图像颜色的四个方面：亮度、对比度、饱和度和色调
                      transforms.ColorJitter( brightness=0.5, contrast=0.5, saturation=0.5, hue=0.5) 

        ])
    def __len__(self) -> int:
        return self.size

    def __getitem__(self, idx: int) -> TypeRotItem:
        img_ts = self.imgseq[idx] #tensor
        img_ts = self.transforms(img_ts)
        index_ts: Tensor = self.indices[idx] # 旋转44
        img_ts = from_img(img_ts, index_ts.item() / self.cls_num, self.target_size)
        img_ts = self.norm(img_ts)

        return img_ts, index_ts

tensor类型图片保存方法

from torchvision import transforms
toPIL= transforms.ToPILImage()  # 这个函数可以将张量转为PIL图片，由小数转为0-255之间的像素值
pic = toPIL(img_ts2)
pic.save('img_ts2.jpg')

训练过程和结果

结果：
一开始使用原始底图110张图片 准确率有80%（感谢simple ocr项目拥有者提供图片），数据增强后可以达到85%准确率，训练集loss可以降低到1.2，python test_RotNetR.py 计算平均误差度数的结果达到1度以内

后续训练， 增加训练图片到300张（保证图片质量较高，多样性丰富），准确度可以达到90%以上

数据预处理

找出相似的图片，先去重

 
import numpy as np
import os
import cv2
from tqdm import tqdm
def ssim(y_true , y_pred):
    u_true = np.mean(y_true)
    u_pred = np.mean(y_pred)
    var_true = np.var(y_true)
    var_pred = np.var(y_pred)
    std_true = np.sqrt(var_true)
    std_pred = np.sqrt(var_pred)
    R = 255
    c1 = np.square(0.01*R)
    c2 = np.square(0.03*R)
    ssim = (2 * u_true * u_pred + c1) * (2 * std_pred * std_true + c2)
    denom = (u_true ** 2 + u_pred ** 2 + c1) * (var_pred + var_true + c2)
    return ssim / denom

def show(image1,image2=''):
    # 创建一个窗口并显示合并后的图片
    if image2=='':
        combined_image =image1
    else:
        combined_image = cv2.hconcat([image1, image2])
    cv2.namedWindow('Combined Image', cv2.WINDOW_NORMAL)
    cv2.imshow('Combined Image', combined_image)

    # 等待用户按下任意键，然后关闭窗口
    cv2.waitKey(0)
    cv2.destroyAllWindows()

# 获取当前文件夹中的所有图片文件路径
image_folder = r'xxx'  # 存放图片的文件夹路径

file_list = os.listdir(image_folder)
result_file_list = []
images = sorted([os.path.join(image_folder, file) for file in file_list if file.endswith(('jpg', 'png', 'jpeg'))])
doubt_list = []

for i in range(len(images)):
    print(i)
    r_list = []
    max_r = 0
    max_j = 0
    img1 = cv2.imread(images[i] )
    # 灰度图像处理
    gray_img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    # _, b_img1 = cv2.threshold(gray_img1, 200, 255, cv2.THRESH_BINARY)
    # ret, img1 = cv2.threshold(img1, 127, 255, cv2.THRESH_BINARY)
    for j in range(i+1,len(images)):
        flag = 0
        try:

            img2 = cv2.imread(images[j])
            gray_img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
            # _, b_img2 = cv2.threshold(gray_img2, 180, 255, cv2.THRESH_BINARY)
            r = ssim(gray_img1, gray_img2)
            if r > 0.99 and r > max_r:
                max_r = r
                max_j = j
        except:
            print('图片已经删除')
    if max_j!=0 and max_r>0.99:
        # show(b_img1, b_img2)
        result_file_list.append(images[i])
        cv2.imwrite('%s_first.jpg'%(i),img1)
        cv2.imwrite('%s_%s_second.jpg'%(i,max_j), cv2.imread(images[max_j]) )
        # show(img1, cv2.imread(images[max_j]))
        #
        # if flag == 1:
        #     doubt_list.append([file_list[i], file_list[j]])
print(result_file_list)#类似('101_164.png', '126_216.png') ('130_112.png', '99_172.png')

修正图片角度

注意旋转图片会造成图片质量降低，使用cv2.INTER_CUBIC填充，之后还可以使用高分辨率工具恢复图像

from flask import Flask, render_template, request
from PIL import Image
import os
import cv2
import numpy as np
from PIL import Image, ImageOps
app = Flask(__name__)

# 获取当前文件夹中的所有图片文件路径
image_folder = 'static/images'  # 存放图片的文件夹路径

images_path = sorted([os.path.join(image_folder, file) for file in os.listdir(image_folder) if file.endswith(('jpg', 'png', 'jpeg'))])
current_index = 0  # 当前显示的图片索引

def get_current_index_by_name(file_name):
    global images_path
    result_index_list = [i for i in range(len(images_path)) if os.path.basename(images_path[i])==file_name]
    return result_index_list[0]


def rotate_image( image, angle, if_fill_white = True):
    '''
    顺时针旋转
    Args:
        image:
        angle:
        if_fill_white:旋转产生的黑边填充为白色

    Returns:

    '''
    # dividing height and width by 2 to get the center of the image
    height, width = image.shape[:2]
    # get the center coordinates of the image to create the 2D rotation matrix
    center = (width / 2, height / 2)

    # using cv2.getRotationMatrix2D() to get the rotation matrix
    rotate_matrix = cv2.getRotationMatrix2D(center=center, angle=angle, scale=1)

    # rotate the image using cv2.warpAffine
    if not if_fill_white:
        rotated_image = cv2.warpAffine(src=image, M=rotate_matrix, dsize=(width, height), flags=cv2.INTER_CUBIC)
    else:
        color = (255, 255) if len(image.shape)==2 else (255, 255,255)
        rotated_image = cv2.warpAffine(src=image, M=rotate_matrix, dsize=(width, height), borderValue=color, flags=cv2.INTER_CUBIC)
    return rotated_image

# 显示当前图片
@app.route('/')
def index():
    global current_index
    img_path = images_path[current_index]
    return render_template('index.html', img_path=img_path)

def cv_imread(file_path):
    cv_img = cv2.imdecode(np.fromfile(file_path,dtype=np.uint8),-1)
    return cv_img

def cv_imwrite(img_path, img):
    cv2.imencode('.jpg', img)[1].tofile(img_path)

# 处理图片切换和旋转请求
@app.route('/action', methods=['POST'])
def handle_action():
    global current_index,images_path
    img_path = images_path[current_index] #原图路径
    action = request.form['action']
    degree = 1 if request.form['input_text']=='' else float(request.form['input_text'])

    # rotated_img = Image.open(img_path)
    rotated_img = cv_imread(img_path)
    if action == 'left':
        rotated_img = rotate_image(rotated_img, degree)
        cv_imwrite(img_path,rotated_img)
        # rotated_img = rotated_img.rotate(1, expand=False, fillcolor='white')
        # rotated_img.save(img_path, quality=95)

    elif action == 'right':
        rotated_img = rotate_image(rotated_img, -degree)
        cv_imwrite(img_path, rotated_img)
        # rotated_img = rotated_img.rotate(-1, expand=False, fillcolor='white')
        # rotated_img.save(img_path, quality = 95)
    else:
        if action == 'next':
            current_index = (current_index + 1) % len(images_path)
        elif action == 'previous':
            current_index = (current_index - 1) % len(images_path)

    return render_template('index.html', img_path= images_path[current_index] )

if __name__ == '__main__':
    app.run(debug=True)

index.html

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Image Viewer</title>
    <style>
        .container {
            text-align: center;
            position: relative; /* 设置相对定位 */
        }

        .image-container {
            display: inline-block;
            text-align: left;
            position: relative; /* 设置相对定位 */
        }

        .image-container img {
            max-width: 50%; /* 调整图片最大宽度为 80% */
            height: auto;
            margin: 10px;
        }

        .button-container {
            margin-top: 20px;
        }
         /* 模拟九宫格水平仪 */
        .nine-grid {
            position: absolute;
            top: 50%;
            left: 50%;
            transform: translate(-50%, -50%);
            width: 80%; /* 九宫格宽度 */
            height: 80%; /* 九宫格高度 */
            border: 1px dashed red; /* 红色虚线边框 */
        }

        /* 水平线 */
        .nine-grid::before,
        .nine-grid::after,
        .nine-grid::first-line {
            content: "";
            position: absolute;
            background-color: red; /* 红色实线 */
            width: 100%;
            height: 1px; /* 水平线粗细 */
        }

        /* 垂直线 */
        .nine-grid::before,
        .nine-grid::after,
        .nine-grid::first-line {
            content: "";
            position: absolute;
            background-color: red; /* 红色实线 */
            width: 1px; /* 垂直线粗细 */
            height: 100%;
        }

        /* 第一条水平线位置 */
        .nine-grid::before {
            top: 50%;
            left: 0;
        }

        /* 第二条水平线位置 */
        .nine-grid::first-line {
            top: 60%;
            left: 0;
        }

        /* 第三条水平线位置 */
        .nine-grid::after {
            top: 70%;
            left: 0;
        }

        /* 第一条垂直线位置 */
        .nine-grid::before {
            top: 0;
            left: 10%;
        }

        /* 第二条垂直线位置 */
        .nine-grid::first-line {
            top: 0;
            left: 25%;
        }

        /* 第三条垂直线位置 */
        .nine-grid::after {
            top: 0;
            left: 40%;
        }


    </style>
</head>
<body>
<div class="container">
    <div class="image-container">
        <div class="nine-grid"></div> <!-- 模拟横竖虚线水平仪 -->
        <img src="{{ img_path }}" alt="Current Image">
    </div>
    {{ img_path }}
    <form action="/action" method="post">
        <div class="button-container">
            <button type="submit" name="action" value="previous">Previous Pic</button>
             <input type="text" name="input_text">
            <button type="submit" name="action" value="left">Left</button>
            <button type="submit" name="action" value="right">Right</button>
            <button type="submit" name="action" value="next">Next Pic</button>
        </div>
    </form>
</div>
</body>
</html>