1 Star 1 Fork 1

guojy/pytorch-captcha

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
该仓库未声明开源许可证文件(LICENSE),使用请关注具体项目描述及其代码上游依赖。
克隆/下载
predict.py 1.52 KB
一键复制 编辑 原始数据 按行查看 历史
ice-tong 提交于 2019-03-23 15:02 +08:00 . Add files via upload
# -*- coding: utf-8 -*-
"""
Created on Wed Feb 13 20:07:17 2019
@author: icetong
"""
import torch
import torch.nn as nn
from models import CNN
from datasets import CaptchaData
from torchvision.transforms import Compose, ToTensor
import matplotlib.pyplot as plot
model_path = './checkpoints/model.pth'
source = [str(i) for i in range(0, 10)]
source += [chr(i) for i in range(97, 97+26)]
alphabet = ''.join(source)
def predict(img_dir='./data/test'):
transforms = Compose([ToTensor()])
dataset = CaptchaData(img_dir, transform=transforms)
cnn = CNN()
if torch.cuda.is_available():
cnn = cnn.cuda()
cnn.eval()
cnn.load_state_dict(torch.load(model_path))
for k, (img, target) in enumerate(dataset):
img = img.view(1, 3, 100, 180).cuda()
target = target.view(1, 4*36).cuda()
output = cnn(img)
output = output.view(-1, 36)
target = target.view(-1, 36)
output = nn.functional.softmax(output, dim=1)
output = torch.argmax(output, dim=1)
target = torch.argmax(target, dim=1)
output = output.view(-1, 4)[0]
target = target.view(-1, 4)[0]
print('pred: '+''.join([alphabet[i] for i in output.cpu().numpy()]))
print('true: '+ ''.join([alphabet[i] for i in target.cpu().numpy()]))
plot.imshow(img.permute((0, 2, 3, 1))[0].cpu().numpy())
plot.show()
if k >= 10: break
if __name__=="__main__":
predict()
Loading...
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
1
https://gitee.com/covwo/pytorch-captcha.git
[email protected]:covwo/pytorch-captcha.git
covwo
pytorch-captcha
pytorch-captcha
master

搜索帮助