INNER CODE UNIT · Python
SOS
dailenson/SDT · test.py:72
SOS = torch.tensor(bs * [[0, 0, 1, 0, 0]]).unsqueeze(1).to(preds)
preds = torch.cat((SOS, preds), 1) # add the SOS token like GT
preds = preds.detach().cpu().numpy()
test_cache = {}
coords = coords.detach().cpu().numpy()
if opt.store_type == 'online':
for i, pred in enumerate(preds):
pred, _ = dxdynp_to_list(preds[i])
coord, _ = dxdynp_to_list(coords[i])
data = {'coordinates': pred, 'writer_id': writer_id[i].item(),
'character_id': character_id[i].item(), 'coords_gt':coord}
data_byte = pickle.dumps(data)
data_id = str(num_count).encode('utf-8')
test_cache[data_id] = data_byte
num_count += 1
test_cache['num_sample'.encode('utf-8')] = str(num_count).encode()
writeCache(test_env, test_cache)