# captcha.py
from model import CNN_sjtu
import torch.nn
import torch
from PIL import Image
from torchvision.transforms import Compose, ToTensor, Resize,Normalize
import time
from selenium.webdriver.common.by import By
from threading import Thread
import warnings

#pytorch版本关系，会有warning，实际上可以忽略
warnings.filterwarnings("ignore")

#定义验证码信息参数
numchar = 4
alphabet = 'abcdefghijklmnopqrstuvwxyz'
width = 100
height = 40

#模型初始化
model_net = CNN_sjtu(num_class=len(alphabet), num_char=int(numchar), width=width, height=height)
model_net = model_net.cuda()
model_net.eval()

def load_net():#load参数 比较慢，用一个额外的线程先load
    global model_net
    print("load net......")
    model_net.load_state_dict(torch.load(r'C:\Users\14242\PycharmProjects'
                                         r'\DL\Pytorch_project\captcha-CNN-验证码识别\weights/model_sjtu_4.path'))#参数位置

thread_load= Thread(target=load_net)
thread_load.start()

def decode_captcha(img_path):
    global model_net
    with torch.no_grad():
        img=Image.open(img_path)
        img = img.convert('RGB')
        transforms = Compose([Resize((height, width)), ToTensor(), Normalize(0, 1)])
        img = transforms(img)
        img = img.view(1, 3, height, width).cuda()
        output = model_net(img)
        output = output.view(-1, len(alphabet))
        output = torch.nn.functional.softmax(output, dim=1)
        output = torch.argmax(output, dim=1)
        output = output.view(-1, numchar)[0]
        return ''.join([alphabet[i] for i in output.cpu().detach().numpy()])

def get_snap(driver):  # 对目标网页进行截屏。这里截的是全屏
    driver.save_screenshot('full_snap.png')
    page_snap_obj=Image.open('full_snap.png')
    return page_snap_obj

def get_image(driver): # 对验证码所在位置进行定位，然后截取验证码图片
    scaling_ratio=1#系统显示的缩放比例
    img = driver.find_element(By.XPATH, '//*[@id="captcha-img"]')#获取验证码在网页中的位置信息
    time.sleep(0.10)#等待一会
    location = img.location#位置
    size = img.size#大小
    left = location['x']*scaling_ratio
    top = location['y']*scaling_ratio
    right = left + size['width']*scaling_ratio
    bottom = top + size['height']*scaling_ratio
    page_snap_obj = get_snap(driver)#获取整个页面截图
    image_obj = page_snap_obj.crop((left, top, right, bottom))#截下验证码
    image_obj.save('captcha.png')

# login.py
from captcha import decode_captcha as decode_c
from captcha import get_image
from selenium import webdriver
import time
import os
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By

#调用谷歌driver的基本操作
options = Options()
driver = webdriver.Chrome(options=options)#调用当前路径下的chromedriver.exe
driver.maximize_window()#最大化

#截图的保存地址
captcha_path='./captcha.png'
snap_path='./full_snap.png'

#两个地址链接
url = 'https://i.sjtu.edu.cn/jaccountlogin'
url ='https://oc.sjtu.edu.cn/login/openid_connect'

driver.get(url) #打开网页

cur_title=driver.title
while(driver.title==cur_title):#失败则一直试，因为验证码可能错误，两次应该就能成功
    time.sleep(0.10)  # 加载等待
    get_image(driver)
    print("推理验证码......")
    captcha_res=decode_c(captcha_path)
    driver.find_element(By.NAME, 'user').send_keys('username')  # 填入用户名
    driver.find_element(By.NAME, 'pass').send_keys('password')  # 填入密码
    driver.find_element(By.NAME,'captcha').send_keys(captcha_res)  # 填入验证码
    driver.find_element(By.ID,"submit-button").click()

os.remove(captcha_path)
os.remove(snap_path)
print("ending......")

#  setting.py 
width_sjtu=100	#图片尺寸
height_sjtu=40
alphabet_sjtu='abcdefghijklmnopqrstuvwxyz'#全是小写字母
#验证码长度
numchar=4
#train:
#  遍历数据集训练的次数
max_epoch=100
#  批处理数量
batch_size=128
#  学习率
base_lr=0.0003

#  训练数据存放路径
train_data_path_sjtu='./captcha-sjtu/train'
#  测试数据存放路径
test_data_path_sjtu= './captcha-sjtu/test'
#  预测数据
samples_path_sjtu = './captcha-sjtu/predict'
#  是否使用gpu
use_gpu= True
#  gpu并行处理进程数
num_workers= 0
#  训练后的模型输出的路径
model_path='./weights'

# datasets.py
import os
from PIL import Image
import torch
from torch.utils.data import Dataset

def img_loader(img_path):
    img = Image.open(img_path)
    # 将图像转换为 RGB
    return img.convert('RGB')

# 处理数据集所在文件夹下的数据
def make_dataset(data_path, alphabet, num_class, num_char):
    # 获取数据集所在文件夹的所有文件名
    img_names = os.listdir(data_path)
    samples = []
    for img_name in img_names:
        # 拼接每个图像数据集的路径
        img_path = os.path.join(data_path, img_name)
        # 找出该图像的label
        target_str = img_name.split('.')[0]
        # 判断lable和结果的长度是否一致
        if len(target_str) != num_char:
            continue

        target = []
        # 创建每个数据的target数组 4 * alphabet，这里使用one hot
        for char in target_str:
            #------如果只看小写要映射成小写-------------------
            # if ord(char)>=65 or ord(char)<=90:
            #     char=chr(ord(char)+32)
            #---------------------------------

            vec = [0] * num_class
            vec[alphabet.find(char)] = 1
            target += vec#要四个数组，如果写在一个数组内，重复的表示不出来

        # 加入数据集
        samples.append((img_path, target))
    # 返回数据集
    return samples


class CaptchaData(Dataset):
    def __init__(self, data_path, num_class=62, num_char=4, transform=None, target_transform=None,
                 alphabet="0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"):
        super(Dataset, self).__init__();
        self.data_path = data_path
        self.num_class = num_class
        self.num_char = num_char
        self.transform = transform
        self.target_transform = target_transform
        self.alphabet = alphabet
        self.samples = make_dataset(self.data_path, self.alphabet,
                                    self.num_class, self.num_char)

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        img_path, target = self.samples[index]
        img = img_loader(img_path)
        # 如果有传入预处理函数，就预处理数据集
        if self.transform is not None:
            img = self.transform(img)
        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, torch.Tensor(target)

#  model.py
import torch.nn as nn
class CNN_sjtu(nn.Module):#4长度和5长度准确率都在96%左右，验证都是正确的
    def __init__(self, num_class=36, num_char=4, width=100, height=40):
        super(CNN_sjtu, self).__init__()
        self.num_class = num_class
        self.num_char = num_char
        # 卷积层后，全连接层的一维数组输入长度
        # 512是卷积处理后图片的通道数，长度和宽度各除以16是因为图像经过了四次2*2池化层（MaxPool2d）
        self.line_size = int(512 * (width // 2 // 2 // 2 // 2) * (height // 2 // 2 // 2 // 2))
        self.conv1 = nn.Sequential(
            # 输入的是RGB图像，所以是3通道。
            # 这里设置该层有16个卷积核，所以输出是16通道
            # padding（1，1）表示在图像上下左右各加1行、1列，保证在卷积后图像大小不变
            nn.Conv2d(3, 16, 3, padding=(1, 1)),
            # 池化层，保留图像每2*2片段像素的最大值
            nn.MaxPool2d(2, 2),
            # 对每个通道的图像都归一化，防止梯度爆炸
            nn.BatchNorm2d(16),
            # 激活函数
            nn.ReLU()
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 64, 3, padding=(1, 1)),
            nn.MaxPool2d(2, 2),
            nn.BatchNorm2d(64),
            nn.ReLU()
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(64, 512, 3, padding=(1, 1)),
            nn.MaxPool2d(2, 2),
            nn.BatchNorm2d(512),
            nn.ReLU()
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(512, 512, 3, padding=(1, 1)),
            nn.MaxPool2d(2, 2),
            nn.BatchNorm2d(512),
            nn.ReLU()
        )
        # 全连接层
        self.fc = nn.Sequential(
            nn.Linear(self.line_size, self.line_size),
            # nn.Identity(),
            # 输出应为 验证码长度*字符的分类数
            nn.Linear(self.line_size, self.num_char * self.num_class)
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        # resize输入数组的尺寸，相当于flatten
        x = x.view(-1, self.line_size)
        x = self.fc(x)

        return x
    

#train.py
import torch
from model import *
from datasets import CaptchaData
from torch.utils.data import DataLoader
from torchvision.transforms import Compose, ToTensor, Resize,Normalize
import time
import os
from setting import *

#导入设定参数
alphabet=alphabet_sjtu
width=width_sjtu
height=height_sjtu

#  训练数据存放路径
train_data_path=train_data_path_sjtu
#  测试数据存放路径
test_data_path=test_data_path_sjtu
#  预测数据
samples_path=samples_path_sjtu

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

device =torch.device("cuda")#设置gpu

# 计算准确度
def calculat_acc(output, target):
    output, target = output.view(-1, len(alphabet)), target.view(-1, len(alphabet)) #字母有26个就是26列
    output = nn.functional.softmax(output, dim=1) #缩放到0-1区间，所有元素和为1
    output = torch.argmax(output, dim=1)    #返回每一列得分最高的索引值，说明预测的是这个位置的字母
    target = torch.argmax(target, dim=1)    #然后dim=1这个维度会消失
    output, target = output.view(-1, int(numchar)), target.view(-1, int(numchar))
    correct_list = []
    for i, j in zip(target, output):
        if torch.equal(i, j):   #如果两个列表相等（相同大小和元素）
            correct_list.append(1)
        else:
            correct_list.append(0)
    acc = sum(correct_list) / len(correct_list)
    return acc


def train():
    # 数据shape的预处理，缩放、转tensor，以及图像处理基本都要使用归一化
    transforms = Compose([Resize((height, width)), ToTensor(),Normalize(0, 1)])
    # 创建训练数据集对象
    train_dataset = CaptchaData(train_data_path, num_class=len(alphabet), num_char=int(numchar), transform=transforms, alphabet=alphabet)
    # 初始化DataLoader，之后训练的数据由它按照我们的要求如batch_size等提供
    train_data_loader = DataLoader(train_dataset, batch_size=batch_size, num_workers=num_workers,
                                   shuffle=True, drop_last=False)
    # 创建测试数据集对象
    test_dataset = CaptchaData(test_data_path, num_class=len(alphabet), num_char=int(numchar), transform=transforms, alphabet=alphabet)
    test_data_loader = DataLoader(test_dataset, batch_size=batch_size,
                                  num_workers=num_workers,drop_last=False)

    print("训练数据量：",train_dataset.__len__(),'\t测试数据量：',test_dataset.__len__())

    # 初始化模型
    cnn = CNN_sjtu(num_class=len(alphabet), num_char=int(numchar), width=width, height=height)
    if use_gpu:
        cnn=cnn.to(device)
  
    #----------------损失函数及优化方法---------------------
    # 使用Adam优化方法
    optimizer = torch.optim.Adam(cnn.parameters(), lr=base_lr)
    # 使用多标签分类的损失函数
    criterion = nn.MultiLabelSoftMarginLoss()
    
    #----------------开始迭代训练-------------------------
    # 训练我们指定的epoch次
    print("开始训练...")
    for epoch in range(max_epoch):
        start_ = time.time()
        loss_history = []
        acc_history = []
        # 切换到训练模式
        cnn.train()
        for img, target in train_data_loader:
            if use_gpu:
                img = img.to(device)
                target = target.to(device)
            # 获取神经网络的输出
            output = cnn(img)
            # 计算损失函数
            loss = criterion(output, target)
            # 初始化梯度
            optimizer.zero_grad()
            # 反向传播计算梯度
            loss.backward()
            # 优化参数
            optimizer.step()
            # 计算准确度
            acc = calculat_acc(output, target)
            acc_history.append(float(acc))
            loss_history.append(float(loss))
        print('epoch:{},train_loss: {:.4}|train_acc: {:.4}'.format(
            epoch,
            torch.mean(torch.Tensor(loss_history)),
            torch.mean(torch.Tensor(acc_history)),
        ))

        with torch.no_grad():
            loss_history = []
            acc_history = []
            # 切换到测试模式
            cnn.eval()
            for img, target in test_data_loader:
                if torch.cuda.is_available():
                    img = img.to(device)
                    target = target.to(device)
                output = cnn(img)

                acc = calculat_acc(output, target)
                acc_history.append(float(acc))
            print('test_loss: {:.4}|test_acc: {:.4}'.format(
                torch.mean(torch.Tensor(loss_history)),
                torch.mean(torch.Tensor(acc_history)),
            ))
        print('epoch: {}|time: {:.4f}'.format(epoch, time.time() - start_))
        torch.save(cnn.state_dict(), os.path.join(model_path, "model_{}.path".format(epoch)))#每个epoch保存一次参数


if __name__ == "__main__":
    train()

#predict.py
import torch
from PIL import Image
from model import *
from torchvision.transforms import Compose, ToTensor, Resize,Normalize
import matplotlib.pyplot as plt
import os
import random
from setting import *

#----------参数设定-------------
model_net = CNN_sjtu()
alphabet=alphabet_sjtu
width=width_sjtu
height=height_sjtu
samples_path=samples_path_sjtu+'/4'
numchar=4

# 获取模型
def load_net():
    global model_net
    # 初始化模型
    model_net = CNN_sjtu(num_class=len(alphabet), num_char=int(numchar), width=width, height=height)
    # 读取参数模型
    if use_gpu:
        model_net = model_net.cuda()
        model_net.eval()
        model_net.load_state_dict(torch.load('./weights/model_sjtu_4.path'))    #加载参数模型
    else:
        model_net.eval()
        model_net.load_state_dict(torch.load(model_path, map_location='cpu'))

# 预测验证码
def predict_image(img):
    global model_net
    with torch.no_grad():
        img = img.convert('RGB')
        transforms = Compose([Resize((height, width)), ToTensor(),Normalize(0, 1)]) #图像变换
        img = transforms(img)

        if use_gpu:
            img = img.view(1, 3, height, width).cuda()
        else:
            img = img.view(1, 3, height, width)
        output = model_net(img) #推理

        output = output.view(-1, len(alphabet))
        output = nn.functional.softmax(output, dim=1)
        output = torch.argmax(output, dim=1)
        output = output.view(-1, numchar)[0]
        return ''.join([alphabet[i] for i in output.cpu().detach().numpy()])  #转换成numpy类型需要先从gpu加载到cpu，然后可以获得字母表的索引


if __name__ == "__main__":
    load_net()
    # 枚举数据所在文件夹
    img_names = os.listdir(samples_path)
    random.shuffle(img_names)
    samples = []
    for img_name in img_names:
        # 拼接每个数据的路径
        img_path = os.path.join(samples_path, img_name)
        img = Image.open(img_path)
        v_code = predict_image(img)
        plt.figure()
        plt.title("{}".format(v_code))
        plt.imshow(img)
        plt.show()