34 lines
1.4 KiB
Python
34 lines
1.4 KiB
Python
|
import numpy as np
|
|||
|
|
import captcha_settings
|
|||
|
|
import torch
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 用torch.zeros()函数生成一个4行36列，值全是0的张量。接着循环标签中的各个字符，将字符在captcha_settings.ALL_CHAR_SET_STR中对应的索引获取到，然后将张量中对应位置的0，改成1。最后要返回一个一维的列表，长度是4*36=144
|
|||
|
|
def encode(label):
|
|||
|
|
"""将字符转为独热码"""
|
|||
|
|
cols = len(captcha_settings.ALL_CHAR_SET_STR)
|
|||
|
|
rows = captcha_settings.MAX_CAPTCHA
|
|||
|
|
result = torch.zeros((rows, cols), dtype=float)
|
|||
|
|
for i, char in enumerate(label):
|
|||
|
|
j = captcha_settings.ALL_CHAR_SET_STR.index(char)
|
|||
|
|
result[i, j] = 1.0
|
|||
|
|
return result.view(1, -1)[0]
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 将模型预测的值从一维转成4行36列的二维张量，然后调用torch.argmax()函数寻找每一行最大值（也就是1）的索引。知道索引后就可以从captcha_settings.ALL_CHAR_SET_STR中找到对应的字符
|
|||
|
|
def decode(pred_result):
|
|||
|
|
"""将独热码转为字符"""
|
|||
|
|
pred_result = pred_result.view(-1, len(captcha_settings.ALL_CHAR_SET_STR))
|
|||
|
|
index_list = torch.argmax(pred_result, dim=1)
|
|||
|
|
text = "".join([captcha_settings.ALL_CHAR_SET_STR[i] for i in index_list])
|
|||
|
|
return text
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
label = "ABCD"
|
|||
|
|
one_hot_label = encode(label)
|
|||
|
|
print(one_hot_label)
|
|||
|
|
decoded_label = decode(one_hot_label)
|
|||
|
|
print(decoded_label)
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
main()
|