def predict(img_dir='./data/test'):
transforms = Compose([Resize(height, weight), 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, height, weight).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
def predict(img_dir=testpath):
transforms = Compose([ToTensor()])
dataset = CaptchaData(img_dir, transform=transforms)
# dataset = MyDataSet()
cnn = CNN()
if torch.cuda.is_available():
cnn = cnn.cuda()
cnn.eval()
cnn.load_state_dict(torch.load(model_path))
else:
cnn.eval()
model = torch.load(model_path, map_location='cpu')
cnn.load_state_dict(model)
for k, (img, target) in enumerate(dataset):
#img = img.view(1, 3, 40, 120).cuda()
#target = target.view(1, 4*62).cuda()
#target = target.view(1, 4*62).cuda()
img = img.view(1, 3, 40, 120)
output = cnn(img)
output = output.view(-1, 62)
target = target.view(-1, 62)
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]
a = ''.join([alphabet[i] for i in output.cpu().numpy()])
b = ''.join([alphabet[i] for i in target.cpu().numpy()]) + '.jpg'
#print('ID: ' + ''.join([alphabet[i] for i in target.cpu().numpy()]) + '.jpg')
#print('label: '+''.join([alphabet[i] for i in output.cpu().numpy()]))
dataframe = pd.DataFrame({'ID': b, 'label': a}, index=[0])
dataframe.to_csv('1.csv', index=False, mode='a')
with open('1.csv', 'r') as r:
lines = r.readlines()
with open('2.csv', 'w') as w:
for i in lines:
if 'label' not in i:
w.write(i)
df = pd.read_csv('2.csv', header=None, names=['ID', 'label'])
df.to_csv('3.csv', index=False)
with open('3.csv', 'r') as r:
lines = r.readlines()
with open(result_path, 'w') as w:
for l in lines:
str = l.strip('0')
w.write(str)
f = pd.read_csv(result_path)
print(f)
os.remove('1.csv')
os.remove('2.csv')
os.remove('3.csv')
def captchaByDir(img_dir):
transforms = Compose([ToTensor()])
dataset = CaptchaData(img_dir, transform=transforms)
lable = []
for k, (img, target) in enumerate(dataset):
if torch.cuda.is_available():
img = img.view(1, 3, 32, 120).cuda()
else:
img = img.view(1, 3, 32, 120)
output = cnn(img)
output = output.view(-1, 36)
output = nn.functional.softmax(output, dim=1)
output = torch.argmax(output, dim=1)
output = output.view(-1, 4)[0]
lable.append(''.join([source[i] for i in output.cpu().numpy()]))
return lable
def train():
transforms = Compose([ToTensor()])
train_dataset = CaptchaData('./data/train', transform=transforms)
train_data_loader = DataLoader(train_dataset,
batch_size=batch_size,
num_workers=0,
shuffle=True,
drop_last=True)
test_data = CaptchaData('./data/test', transform=transforms)
test_data_loader = DataLoader(test_data,
batch_size=batch_size,
num_workers=0,
shuffle=True,
drop_last=True)
cnn = CNN()
if torch.cuda.is_available():
cnn.cuda()
if restor:
cnn.load_state_dict(torch.load(model_path))
# freezing_layers = list(cnn.named_parameters())[:10]
# for param in freezing_layers:
# param[1].requires_grad = False
# print('freezing layer:', param[0])
optimizer = torch.optim.Adam(cnn.parameters(), lr=base_lr)
criterion = nn.MultiLabelSoftMarginLoss()
for epoch in range(max_epoch):
start_ = time.time()
loss_history = []
acc_history = []
cnn.train()
for img, target in train_data_loader:
img = Variable(img)
target = Variable(target)
if torch.cuda.is_available():
img = img.cuda()
target = target.cuda()
output = cnn(img)
loss = criterion(output, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
acc = calculat_acc(output, target)
acc_history.append(acc)
loss_history.append(loss)
print('train_loss: {:.4}|train_acc: {:.4}'.format(
torch.mean(torch.Tensor(loss_history)),
torch.mean(torch.Tensor(acc_history)),
))
loss_history = []
acc_history = []
cnn.eval()
for img, target in test_data_loader:
img = Variable(img)
target = Variable(target)
if torch.cuda.is_available():
img = img.cuda()
target = target.cuda()
output = cnn(img)
acc = calculat_acc(output, target)
acc_history.append(acc)
loss_history.append(float(loss))
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(), model_path)
def train():
transforms = Compose([Resize((height, width)), ToTensor()])
train_dataset = CaptchaData(train_data_path, num_class=len(alphabet), num_char=int(numchar), transform=transforms, alphabet=alphabet)
train_data_loader = DataLoader(train_dataset, batch_size=batch_size, num_workers=num_workers,
shuffle=True, drop_last=True)
test_data = CaptchaData(test_data_path, num_class=len(alphabet), num_char=int(numchar), transform=transforms, alphabet=alphabet)
test_data_loader = DataLoader(test_data, batch_size=batch_size,
num_workers=num_workers, shuffle=True, drop_last=True)
cnn = CNN(num_class=len(alphabet), num_char=int(numchar), width=width, height=height)
if use_gpu:
cnn.cuda()
optimizer = torch.optim.Adam(cnn.parameters(), lr=base_lr)
criterion = nn.MultiLabelSoftMarginLoss()
for epoch in range(max_epoch):
start_ = time.time()
loss_history = []
acc_history = []
cnn.train()
for img, target in train_data_loader:
img = Variable(img)
target = Variable(target)
if use_gpu:
img = img.cuda()
target = target.cuda()
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)),
))
loss_history = []
acc_history = []
cnn.eval()
for img, target in test_data_loader:
img = Variable(img)
target = Variable(target)
if torch.cuda.is_available():
img = img.cuda()
target = target.cuda()
output = cnn(img)
acc = calculat_acc(output, target)
acc_history.append(float(acc))
loss_history.append(float(loss))
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)))