def save(data, target, target_affinity, no):
"""
Saving the synthesised outputs in between the training
:param data: image as tensor
:param target: character heatmap target as tensor
:param target_affinity: affinity heatmap target as tensor
:param no: current iteration number
:return: None
"""
data = data.data.cpu().numpy()
target = target.data.cpu().numpy()
target_affinity = target_affinity.data.cpu().numpy()
base = './debug/' + str(no) + '/'
os.makedirs(base, exist_ok=True)
for i in range(1):
os.makedirs(base + str(i), exist_ok=True)
plt.imsave(base + str(i) + '/image.png',
denormalize_mean_variance(data[i].transpose(1, 2, 0)))
plt.imsave(base + str(i) + '/target_characters.png',
target[i, :, :],
cmap='gray')
plt.imsave(base + str(i) + '/target_affinity.png',
target_affinity[i, :, :],
cmap='gray')
def save(no, dataset_name, output, image, character_map, affinity_map, character_weight, affinity_weight):
os.makedirs('Temporary/' + str(no), exist_ok=True)
for __, _ in enumerate(dataset_name):
os.makedirs('Temporary/'+str(no)+'/'+str(__), exist_ok=True)
generated = generate_word_bbox(
output[__, 0].data.cpu().numpy(), output[__, 1].data.cpu().numpy(),
config.threshold_character, config.threshold_affinity, config.threshold_word,
config.threshold_character_upper, config.threshold_affinity_upper, config.scale_character, config.scale_affinity
)
output_image = denormalize_mean_variance(
image[__].data.cpu().numpy().transpose(1, 2, 0))
cv2.drawContours(
output_image, generated['word_bbox'], -1, (0, 255, 0), 2)
plt.imsave('Temporary/'+str(no)+'/'+str(__) +
'/image_.png', output_image)
cv2.imwrite('Temporary/'+str(no)+'/'+str(__)+'/char_map.png',
np.uint8(output[__, 0].data.cpu().numpy()*255))
cv2.imwrite('Temporary/'+str(no)+'/'+str(__)+'/aff_map.png',
np.uint8(output[__, 1].data.cpu().numpy()*255))
cv2.imwrite(
'Temporary/' + str(no) + '/' + str(__) +
'/char_map_threshold_upper.png',
np.uint8(np.float32(output[__, 0].data.cpu().numpy() > config.threshold_character_upper) * 255))
cv2.imwrite(
'Temporary/' + str(no) + '/' + str(__) +
'/aff_map_threshold_upper.png',
np.uint8(np.float32(output[__, 1].data.cpu().numpy() > config.threshold_affinity_upper) * 255))
cv2.imwrite(
'Temporary/' + str(no) + '/' + str(__) +
'/char_map_threshold_lower.png',
np.uint8(np.float32(output[__, 0].data.cpu().numpy() > config.threshold_character) * 255))
cv2.imwrite(
'Temporary/' + str(no) + '/' + str(__) +
'/aff_map_threshold_lower.png',
np.uint8(np.float32(output[__, 1].data.cpu().numpy() > config.threshold_affinity) * 255))
cv2.imwrite(
'Temporary/'+str(no)+'/'+str(__)+'/target_char_map.png', np.uint8(character_map[__].data.cpu().numpy()*255))
cv2.imwrite(
'Temporary/'+str(no)+'/'+str(__)+'/target_affinity_map.png', np.uint8(affinity_map[__].data.cpu().numpy()*255))
cv2.imwrite(
'Temporary/'+str(no)+'/'+str(__)+'/weight_char_map.png', np.uint8(character_weight[__].data.cpu().numpy()*255))
cv2.imwrite(
'Temporary/'+str(no)+'/'+str(__)+'/weight_affinity_map.png', np.uint8(affinity_weight[__].data.cpu().numpy()*255))
def save(no, dataset_name, output, image, character_map, affinity_map, character_weight, affinity_weight):
os.makedirs('Temporary/' + str(no), exist_ok=True)
for __, _ in enumerate(dataset_name):
os.makedirs('Temporary/'+str(no)+'/'+str(__), exist_ok=True)
plt.imsave('Temporary/'+str(no)+'/'+str(__)+'/image_.png', denormalize_mean_variance(
image[__].data.cpu().numpy().transpose(1, 2, 0)))
plt.imsave(
'Temporary/'+str(no)+'/'+str(__)+'/char_map.png', output[__, 0].data.cpu().numpy(),
cmap='gray')
plt.imsave(
'Temporary/'+str(no)+'/'+str(__)+'/aff_map.png', output[__, 1].data.cpu().numpy(),
cmap='gray')
plt.imsave(
'Temporary/'+str(no)+'/'+str(__)+'/target_char_map.png', character_map[__].data.cpu().numpy(),
cmap='gray')
plt.imsave(
'Temporary/'+str(no)+'/'+str(__)+'/target_affinity_map.png', affinity_map[__].data.cpu().numpy(),
cmap='gray')
plt.imsave(
'Temporary/'+str(no)+'/'+str(__)+'/weight_char_map.png', character_weight[__].data.cpu().numpy(),
cmap='gray')
plt.imsave(
'Temporary/'+str(no)+'/'+str(__)+'/weight_affinity_map.png', affinity_weight[__].data.cpu().numpy(),
cmap='gray')
def save(data, output, target, target_affinity, drawn_image, no):
"""
Saving the synthesised outputs in between the training
:param data: image as tensor
:param output: predicted output from the model as tensor
:param target: character heatmap target as tensor
:param target_affinity: affinity heatmap target as tensor
:param no: current iteration number
:return: None
"""
output = output.data.cpu().numpy()
data = data.data.cpu().numpy()
target = target.data.cpu().numpy()
target_affinity = target_affinity.data.cpu().numpy()
drawn_image = drawn_image.data.cpu().numpy()
batch_size = output.shape[0]
base = config.DataLoader_JPN_SYNTH_Train_Synthesis + str(no) + '/'
os.makedirs(base, exist_ok=True)
for i in range(batch_size):
os.makedirs(base + str(i), exist_ok=True)
character_bbox = output[i, 0, :, :]
affinity_bbox = output[i, 1, :, :]
plt.imsave(base + str(i) + '/image.png',
denormalize_mean_variance(data[i].transpose(1, 2, 0)))
plt.imsave(base + str(i) + '/target_characters.png', target[i, :, :])
plt.imsave(base + str(i) + '/target_affinity.png',
target_affinity[i, :, :])
blob = np.logical_or(
target[i, :, :] > config.threshold_character,
target_affinity[i, :, :] > config.threshold_affinity)
blob = np.float32(blob)
plt.imsave(base + str(i) + '/blob.png', blob)
plt.imsave(base + str(i) + '/pred_characters.png', character_bbox)
plt.imsave(base + str(i) + '/pred_affinity.png', affinity_bbox)
# Thresholding the character and affinity heatmap
plt.imsave(base + str(i) + '/pred_characters_thresh.png',
np.float32(character_bbox > config.threshold_character))
plt.imsave(base + str(i) + '/pred_affinity_thresh.png',
np.float32(affinity_bbox > config.threshold_affinity))
plt.imsave(base + str(i) + '/drawn_image.png', drawn_image[i])
def test(model, iteration):
"""
Test the weak-supervised model
:param model: Pre-trained model on SynthText
:param iteration: Iteration Number
:return: F-score, loss
"""
os.makedirs(config.save_path + '/Test_'+str(iteration), exist_ok=True)
dataloader = DataLoader(
DataLoaderEvalOther('test'),
batch_size=config.batch_size['test'],
num_workers=config.num_workers['test'],
shuffle=False, worker_init_fn=_init_fn
)
true_positive = 0
false_positive = 0
num_positive = 0
with torch.no_grad():
model.eval()
iterator = tqdm(dataloader)
all_accuracy = []
ground_truth = dataloader.dataset.gt
for no, (image, image_name, original_dim, item) in enumerate(iterator):
annots = []
for i in item:
annot = ground_truth['annots'][dataloader.dataset.imnames[i]]
annots.append(annot)
if config.use_cuda:
image = image.cuda()
output = model(image)
if type(output) == list:
output = torch.cat(output, dim=0)
output = output.data.cpu().numpy()
output[output > 1] = 1
output[output < 0] = 0
original_dim = original_dim.cpu().numpy()
f_score = []
for i in range(output.shape[0]):
# --------- Resizing it back to the original image size and saving it ----------- #
cur_image = denormalize_mean_variance(image[i].data.cpu().numpy().transpose(1, 2, 0))
max_dim = original_dim[i].max()
resizing_factor = 768 / max_dim
before_pad_dim = [int(original_dim[i][0] * resizing_factor), int(original_dim[i][1] * resizing_factor)]
height_pad = (768 - before_pad_dim[0]) // 2
width_pad = (768 - before_pad_dim[1]) // 2
cur_image = cv2.resize(
cur_image[height_pad:height_pad + before_pad_dim[0], width_pad:width_pad + before_pad_dim[1]],
(original_dim[i][1], original_dim[i][0]))
cv2.drawContours(cur_image, resize_bbox(original_dim[i], output[i], config)['word_bbox'], -1, (0, 255, 0), 2)
cv2.drawContours(cur_image, np.array(annots[i]['bbox']), -1, (0, 0, 255), 2)
plt.imsave(
config.save_path + '/Test_' + str(iteration) + '/' + image_name[i],
cur_image.astype(np.uint8))
score_calc = calculate_fscore(
resize_bbox(original_dim[i], output[i], config)['word_bbox'][:, :, 0, :],
np.array(annots[i]['bbox']),
text_target=annots[i]['text'],
)
f_score.append(
score_calc['f_score']
)
true_positive += score_calc['true_positive']
false_positive += score_calc['false_positive']
num_positive += score_calc['num_positive']
# --------------- PostProcessing for creating the targets for the next iteration ---------------- #
all_accuracy.append(np.mean(f_score))
precision = true_positive / (true_positive + false_positive)
recall = true_positive / num_positive
iterator.set_description(
'F-score: ' + str(np.mean(all_accuracy)) + '| Cumulative F-score: '
+ str(2*precision*recall/(precision + recall)))
torch.cuda.empty_cache()
return 2*precision*recall/(precision + recall), precision, recall
def synthesize(dataloader, model, base_path_affinity, base_path_character,
base_path_bbox, base_path_char, base_path_aff, base_path_json):
"""
Given a path to a set of images, and path to a pre-trained model, generate the character heatmap and affinity heatmap
:param dataloader: A Pytorch dataloader for loading and resizing the images of the folder
:param model: A pre-trained model
:param base_path_affinity: Path where to store the predicted affinity heatmap
:param base_path_character: Path where to store the predicted character heatmap
:param base_path_bbox: Path where to store the word_bbox overlapped on images
:param base_path_aff: Path where to store the predicted affinity bbox
:param base_path_char: Path where to store the predicted character bbox
:param base_path_json: Path where to store the predicted bbox in json format
:return: None
"""
with torch.no_grad():
model.eval()
iterator = tqdm(dataloader)
for no, (image, image_name, original_dim) in enumerate(iterator):
if config.use_cuda:
image = image.cuda()
output = model(image)
if type(output) == list:
# If using custom DataParallelModel this is necessary to convert the list to tensor
output = torch.cat(output, dim=0)
output = output.data.cpu().numpy()
output[output < 0] = 0
output[output > 1] = 1
original_dim = original_dim.cpu().numpy()
for i in range(output.shape[0]):
# --------- Resizing it back to the original image size and saving it ----------- #
image_i = denormalize_mean_variance(
image[i].data.cpu().numpy().transpose(1, 2, 0))
max_dim = original_dim[i].max()
resizing_factor = 768 / max_dim
before_pad_dim = [
int(original_dim[i][0] * resizing_factor),
int(original_dim[i][1] * resizing_factor)
]
output[i, :, :, :] = np.uint8(output[i, :, :, :] * 255)
height_pad = (768 - before_pad_dim[0]) // 2
width_pad = (768 - before_pad_dim[1]) // 2
image_i = cv2.resize(
image_i[height_pad:height_pad + before_pad_dim[0],
width_pad:width_pad + before_pad_dim[1]],
(original_dim[i][1], original_dim[i][0]))
character_bbox = cv2.resize(
output[i, 0, height_pad:height_pad + before_pad_dim[0],
width_pad:width_pad + before_pad_dim[1]],
(original_dim[i][1], original_dim[i][0])) / 255
affinity_bbox = cv2.resize(
output[i, 1, height_pad:height_pad + before_pad_dim[0],
width_pad:width_pad + before_pad_dim[1]],
(original_dim[i][1], original_dim[i][0])) / 255
predicted_bbox = generate_word_bbox(
character_bbox,
affinity_bbox,
character_threshold=config.threshold_character,
affinity_threshold=config.threshold_affinity,
word_threshold=config.threshold_word,
character_threshold_upper=config.threshold_character_upper,
affinity_threshold_upper=config.threshold_affinity_upper,
scaling_character=config.scale_character,
scaling_affinity=config.scale_affinity)
word_bbox = predicted_bbox['word_bbox']
char_bbox = np.concatenate(predicted_bbox['characters'],
axis=0)
aff_bbox = np.concatenate(predicted_bbox['affinity'], axis=0)
word_image = image_i.copy()
char_image = image_i.copy()
aff_image = image_i.copy()
cv2.drawContours(word_image, word_bbox, -1, (0, 255, 0), 2)
cv2.drawContours(char_image, char_bbox, -1, (0, 255, 0), 2)
cv2.drawContours(aff_image, aff_bbox, -1, (0, 255, 0), 2)
plt.imsave(
base_path_char + '/' +
'.'.join(image_name[i].split('.')[:-1]) + '.png',
char_image)
plt.imsave(
base_path_aff + '/' +
'.'.join(image_name[i].split('.')[:-1]) + '.png',
aff_image)
plt.imsave(
base_path_bbox + '/' +
'.'.join(image_name[i].split('.')[:-1]) + '.png',
word_image)
plt.imsave(
base_path_character + '/' +
'.'.join(image_name[i].split('.')[:-1]) + '.png',
np.float32(character_bbox > config.threshold_character),
cmap='gray')
plt.imsave(
base_path_affinity + '/' +
'.'.join(image_name[i].split('.')[:-1]) + '.png',
np.float32(affinity_bbox > config.threshold_affinity),
cmap='gray')
predicted_bbox['word_bbox'] = predicted_bbox[
'word_bbox'].tolist()
predicted_bbox['characters'] = [
_.tolist() for _ in predicted_bbox['characters']
]
predicted_bbox['affinity'] = [
_.tolist() for _ in predicted_bbox['affinity']
]
with open(
base_path_json + '/' +
'.'.join(image_name[i].split('.')[:-1]) + '.json',
'w') as f:
json.dump(predicted_bbox, f)
def generate_next_targets(original_dim, output, image, base_target_path,
image_name, annots, dataloader, no):
if 'datapile' in config.dataset_name:
image_name = image_name.split('/')[-1]
# visualize = config.visualize_generated and no % config.visualize_freq == 0 and no != 0
visualize = config.visualize_generated # Just for debuging
max_dim = original_dim.max()
resizing_factor = 768 / max_dim
before_pad_dim = [
int(original_dim[0] * resizing_factor),
int(original_dim[1] * resizing_factor)
]
output = np.uint8(output * 255)
height_pad = (768 - before_pad_dim[0]) // 2
width_pad = (768 - before_pad_dim[1]) // 2
character_bbox = cv2.resize(
output[0, height_pad:height_pad + before_pad_dim[0],
width_pad:width_pad + before_pad_dim[1]],
(original_dim[1] // 2, original_dim[0] // 2)) / 255
affinity_bbox = cv2.resize(
output[1, height_pad:height_pad + before_pad_dim[0],
width_pad:width_pad + before_pad_dim[1]],
(original_dim[1] // 2, original_dim[0] // 2)) / 255
# Generating word-bbox given character and affinity heatmap
generated_targets = generate_word_bbox(
character_bbox,
affinity_bbox,
character_threshold=config.threshold_character,
affinity_threshold=config.threshold_affinity,
word_threshold=config.threshold_word,
character_threshold_upper=config.threshold_character_upper,
affinity_threshold_upper=config.threshold_affinity_upper,
scaling_character=config.scale_character,
scaling_affinity=config.scale_affinity)
generated_targets['word_bbox'] = generated_targets['word_bbox'] * 2
generated_targets['characters'] = [
i * 2 for i in generated_targets['characters']
]
generated_targets['affinity'] = [
i * 2 for i in generated_targets['affinity']
]
if visualize:
character_bbox = cv2.resize((character_bbox * 255).astype(np.uint8),
(original_dim[1], original_dim[0])) / 255
affinity_bbox = cv2.resize((affinity_bbox * 255).astype(np.uint8),
(original_dim[1], original_dim[0])) / 255
image_i = denormalize_mean_variance(image.data.cpu().numpy().transpose(
1, 2, 0))
image_i = cv2.resize(
image_i[height_pad:height_pad + before_pad_dim[0],
width_pad:width_pad + before_pad_dim[1]],
(original_dim[1], original_dim[0]))
# Saving affinity heat map
plt.imsave(base_target_path + '_predicted/affinity/' +
'.'.join(image_name.split('.')[:-1]) + '.png',
np.float32(affinity_bbox > config.threshold_affinity_upper),
cmap='gray')
# Saving character heat map
plt.imsave(
base_target_path + '_predicted/character/' +
'.'.join(image_name.split('.')[:-1]) + '.png',
np.float32(character_bbox > config.threshold_character_upper),
cmap='gray')
cv2.drawContours(image_i, generated_targets['word_bbox'], -1,
(0, 255, 0), 2)
# Saving word bbox drawn on the original image
plt.imsave(
base_target_path + '_predicted/word_bbox/' +
'.'.join(image_name.split('.')[:-1]) + '.png', image_i)
predicted_word_bbox = generated_targets['word_bbox'].copy()
# --------------- PostProcessing for creating the targets for the next iteration ---------------- #
generated_targets = get_weighted_character_target(
generated_targets, {
'bbox': annots['bbox'],
'text': annots['text']
}, dataloader.dataset.unknown, config.threshold_fscore,
config.weight_threshold)
target_word_bbox = generated_targets['word_bbox'].copy()
f_score = calculate_fscore(
predicted_word_bbox[:, :, 0, :],
target_word_bbox[:, :, 0, :],
text_target=annots['text'],
unknown=dataloader.dataset.gt['unknown'])['f_score']
if visualize:
image_i = denormalize_mean_variance(image.data.cpu().numpy().transpose(
1, 2, 0))
image_i = cv2.resize(
image_i[height_pad:height_pad + before_pad_dim[0],
width_pad:width_pad + before_pad_dim[1]],
(original_dim[1], original_dim[0]))
# Generated word_bbox after postprocessing
cv2.drawContours(image_i, generated_targets['word_bbox'], -1,
(0, 255, 0), 2)
# Saving word bbox after postprocessing
plt.imsave(
base_target_path + '_next_target/word_bbox/' +
'.'.join(image_name.split('.')[:-1]) + '.png', image_i)
# Generate affinity heatmap after postprocessing
affinity_target, affinity_weight_map = generate_target_others(
(image_i.shape[0], image_i.shape[1]),
generated_targets['affinity'].copy(),
np.array(generated_targets['weights'])[:, 1])
# Generate character heatmap after postprocessing
character_target, characters_weight_map = generate_target_others(
(image_i.shape[0], image_i.shape[1]),
generated_targets['characters'].copy(),
np.array(generated_targets['weights'])[:, 0])
# Saving the affinity heatmap
plt.imsave(base_target_path + '_next_target/affinity/' +
'.'.join(image_name.split('.')[:-1]) + '.png',
affinity_target,
cmap='gray')
# Saving the character heatmap
plt.imsave(base_target_path + '_next_target/character/' +
'.'.join(image_name.split('.')[:-1]) + '.png',
character_target,
cmap='gray')
# Saving the affinity weight map
plt.imsave(base_target_path + '_next_target/affinity_weight/' +
'.'.join(image_name.split('.')[:-1]) + '.png',
affinity_weight_map,
cmap='gray')
# Saving the character weight map
plt.imsave(base_target_path + '_next_target/character_weight/' +
'.'.join(image_name.split('.')[:-1]) + '.png',
characters_weight_map,
cmap='gray')
# Saving the target for next iteration in json format
generated_targets['word_bbox'] = generated_targets['word_bbox'].tolist()
generated_targets['characters'] = [
word_i.tolist() for word_i in generated_targets['characters']
]
generated_targets['affinity'] = [
word_i.tolist() for word_i in generated_targets['affinity']
]
with open(base_target_path + '/' + image_name + '.json', 'w') as f:
json.dump(generated_targets, f)
return f_score
def synthesize(
dataloader,
model, base_path_affinity, base_path_character, base_path_bbox, base_path_char, base_path_aff, base_path_json):
"""
Given a path to a set of images, and path to a pre-trained model, generate the character heatmap and affinity heatmap
:param dataloader: A Pytorch dataloader for loading and resizing the images of the folder
:param model: A pre-trained model
:param base_path_affinity: Path where to store the predicted affinity heatmap
:param base_path_character: Path where to store the predicted character heatmap
:param base_path_bbox: Path where to store the word_bbox overlapped on images
:param base_path_aff: Path where to store the predicted affinity bbox
:param base_path_char: Path where to store the predicted character bbox
:param base_path_json: Path where to store the predicted bbox in json format
:return: None
"""
with torch.no_grad():
model.eval()
iterator = tqdm(dataloader)
for no, (image, image_name, original_dim) in enumerate(iterator):
if config.use_cuda:
image = image.cuda()
output = model(image)
if type(output) == list:
# If using custom DataParallelModel this is necessary to convert the list to tensor
output = torch.cat(output, dim=0)
output = output.data.cpu().numpy()
output[output < 0] = 0
output[output > 1] = 1
original_dim = original_dim.cpu().numpy()
for i in range(output.shape[0]):
# --------- Resizing it back to the original image size and saving it ----------- #
image_i = denormalize_mean_variance(image[i].data.cpu().numpy().transpose(1, 2, 0))
max_dim = original_dim[i].max()
resizing_factor = 768/max_dim
before_pad_dim = [int(original_dim[i][0]*resizing_factor), int(original_dim[i][1]*resizing_factor)]
output[i, :, :, :] = np.uint8(output[i, :, :, :]*255)
height_pad = (768 - before_pad_dim[0])//2
width_pad = (768 - before_pad_dim[1])//2
image_i = cv2.resize(
image_i[height_pad:height_pad + before_pad_dim[0], width_pad:width_pad + before_pad_dim[1]],
(original_dim[i][1], original_dim[i][0])
)
character_bbox = cv2.resize(
output[i, 0, height_pad:height_pad + before_pad_dim[0], width_pad:width_pad + before_pad_dim[1]],
(original_dim[i][1], original_dim[i][0])
)/255
affinity_bbox = cv2.resize(
output[i, 1, height_pad:height_pad + before_pad_dim[0], width_pad:width_pad + before_pad_dim[1]],
(original_dim[i][1], original_dim[i][0])
)/255
predicted_bbox = generate_word_bbox(
character_bbox,
affinity_bbox,
character_threshold=config.threshold_character,
affinity_threshold=config.threshold_affinity,
word_threshold=config.threshold_word,
character_threshold_upper=config.threshold_character_upper,
affinity_threshold_upper=config.threshold_affinity_upper,
scaling_character=config.scale_character,
scaling_affinity=config.scale_affinity
)
word_bbox = predicted_bbox['word_bbox']
char_bbox = np.concatenate(predicted_bbox['characters'], axis=0)
aff_bbox = np.concatenate(predicted_bbox['affinity'], axis=0)
word_image = image_i.copy()
char_image = image_i.copy()
aff_image = image_i.copy()
cv2.drawContours(word_image, word_bbox, -1, (0, 255, 0), 2)
cv2.drawContours(char_image, char_bbox, -1, (0, 255, 0), 2)
cv2.drawContours(aff_image, aff_bbox, -1, (0, 255, 0), 2)
# plt.imsave(
# base_path_char + '/' + '.'.join(image_name[i].split('.')[:-1]) + '.png',
# char_image)
# plt.imsave(
# base_path_aff + '/' + '.'.join(image_name[i].split('.')[:-1]) + '.png',
# aff_image)
plt.imsave(
base_path_bbox + '/' + '.'.join(image_name[i].split('.')[:-1]) + '.png',
word_image)
# plt.imsave(
# base_path_character + '/' + '.'.join(image_name[i].split('.')[:-1]) + '.png',
# np.float32(character_bbox > config.threshold_character),
# cmap='gray')
# plt.imsave(
# base_path_affinity+'/'+'.'.join(image_name[i].split('.')[:-1])+'.png',
# np.float32(affinity_bbox > config.threshold_affinity),
# cmap='gray')
predicted_bbox['word_bbox'] = predicted_bbox['word_bbox'].tolist()
predicted_bbox['characters'] = [_.tolist() for _ in predicted_bbox['characters']]
predicted_bbox['affinity'] = [_.tolist() for _ in predicted_bbox['affinity']]
with open(base_path_json + '/' + '.'.join(image_name[i].split('.')[:-1])+'.json', 'w') as f:
json.dump(predicted_bbox, f)
boxes_printed = 0
BOX = []
for boxes in predicted_bbox['word_bbox']:
for box in boxes:
BOX.append(box[0])
BOX1= []
BOX2 = []
count = 1
for boxes in BOX:
BOX1.append(boxes)
if count%4 ==0:
BOX2.append(BOX1)
BOX1 = []
count += 1
for box in BOX2:
x = [p[0] for p in box]
y = [p[1] for p in box]
min_x = min(x)
max_x = max(x)
min_y = min(y)
max_y = max(y)
img = image_i.copy()
crop_img = img[min_y:max_y, min_x:max_x]
text = pytesseract.image_to_string(crop_img).rstrip()
boxes_printed += 1
f = open("./text_folder" + '/' + '.'.join(image_name[i].split('.')[:-1]) + '.txt', "a")
f.write(text)
f.write('\n')
f.close()
f = open("./text_folder" + '/' + '.'.join(image_name[i].split('.')[:-1]) + '.txt', "a")
f.write("Total number of boxes : " + str(len(BOX2)))
f.write('\n')
f.write("Total number of boxes printed: " + str(boxes_printed))
f.close()