def model_test(line_img):
img = cv.imread(line_img, 0)
img = cv.bitwise_not(img)
img = np.repeat(img[..., np.newaxis], 3, -1)
img = pp.normalization(img)
prep = T.Compose([
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
test_tensor = prep(img)
test_batch = test_tensor.unsqueeze(0)
modelito = model.make_model(2)
h = modelito.forward(test_batch)
tensor_2_numpy = np.transpose(h.detach().numpy())
return 0
def train(device, source_path, charset_base, max_text_length, batch_size, lr, epochs, target_path):
device = torch.device(device)
model = make_model(tokenizer.vocab_size, hidden_dim=512, nheads=8,
num_encoder_layers=6, num_decoder_layers=6)
model.to(device)
transform = T.Compose([
T.Resize((100, 1300)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
train_loader = torch.utils.data.DataLoader(
DataGenerator(source_path, charset_base, max_text_length, 'train', transform), batch_size=batch_size,
shuffle=False, num_workers=2)
val_loader = torch.utils.data.DataLoader(
DataGenerator(source_path, charset_base, max_text_length, 'valid', transform), batch_size=batch_size,
shuffle=False, num_workers=2)
criterion = LabelSmoothing(size=tokenizer.vocab_size, padding_idx=0, smoothing=0.1)
criterion.to(device) # learning rate
optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=.0004)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1, gamma=0.95)
run_epochs(model, criterion, optimizer, scheduler, train_loader, val_loader, epochs, tokenizer, target_path, device)
compression_opts=9)
print(f"[OK] {i} partition.")
print(f"Transformation finished.")
elif args.image:
img = pp.preprocess(args.image, input_size=input_size)
#making image compitable with resnet
img = np.repeat(img[..., np.newaxis], 3, -1)
x_test = pp.normalization(img)
model = make_model(tokenizer.vocab_size,
hidden_dim=256,
nheads=4,
num_encoder_layers=4,
num_decoder_layers=4)
device = torch.device(args.device)
model.to(device)
transform = T.Compose([T.ToTensor()])
if os.path.exists(target_path):
model.load_state_dict(torch.load(target_path))
else:
print('No model checkpoint found')
prediction = single_image_inference(model, x_test, tokenizer,
transform, device)
print("\n####################################")