Spaces:
Runtime error
Runtime error
File size: 6,107 Bytes
6163604 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 |
'''
This file contains functions to visualize the heatmap and detected bounding boxes
'''
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import os
import numpy as np
import cv2
def draw_stitched_boxes(im, data, outpath):
# Create figure and axes
fig, ax = plt.subplots(1)
# sort based on the confs. Confs is column 4
data = data[data[:, 4].argsort()]
# Display the image
ax.imshow(im)
width, height, channels = im.shape
heatmap = np.zeros([width, height])
for box in data:
heatmap[int(box[1]):int(box[3]), int(box[0]):int(box[2])] = box[4]
# Following line makes sure that all the heatmaps are in the scale, 0 to 1
# So color assigned to different scores are consistent across heatmaps for
# different images
heatmap[0:1, 0:1] = 1
heatmap[0:1, 1:2] = 0
plt.imshow(heatmap, alpha=0.4, cmap='hot', interpolation='nearest')
plt.colorbar()
plt.title("Stitching visualization")
plt.show()
plt.savefig(outpath, dpi=600)
plt.close()
def draw_all_boxes(im, data, recognized_boxes, gt_boxes, outpath):
if len(data) == 0:
return
# Create figure and axes
fig, ax = plt.subplots(1)
# sort based on the confs. Confs is column 4
data = data[data[:, 4].argsort()]
# Display the image
ax.imshow(im)
width, height, channels = im.shape
heatmap = np.zeros([width, height])
if data is not None:
for box in data:
heatmap[int(box[1]):int(box[3]), int(box[0]):int(box[2])] = box[4]
#rect = patches.Rectangle((box[0], box[1]), box[2] - box[0], box[3] - box[1],
# linewidth=0.25, edgecolor='m', facecolor='none')
#Add the patch to the Axes
#ax.add_patch(rect)
if recognized_boxes is not None:
# recognized boxes are green
for box in recognized_boxes:
rect = patches.Rectangle((box[0], box[1]), box[2] - box[0], box[3] - box[1],
linewidth=1, edgecolor='g', facecolor='none')
# Add the patch to the Axes
ax.add_patch(rect)
if gt_boxes is not None:
# ground truth are red
for box in gt_boxes:
rect = patches.Rectangle((box[0], box[1]), box[2] - box[0], box[3] - box[1],
linewidth=0.25, edgecolor='b', facecolor='none')
# Add the patch to the Axes
ax.add_patch(rect)
# Following line makes sure that all the heatmaps are in the scale, 0 to 1
# So color assigned to different scores are consistent across heatmaps for
# different images
heatmap[0:1, 0:1] = 1
heatmap[0:1, 1:2] = 0
plt.imshow(heatmap, alpha=0.4, cmap='hot', interpolation='nearest')
plt.colorbar()
plt.title("Stitching visualization")
plt.show()
plt.savefig(outpath, dpi=600)
plt.close()
def draw_boxes_cv(image, recognized_boxes, gt_boxes, outpath):
'''
:param image
:param recognized_boxes
:param outpath: save as outpath. Should be complete image path with extension
:return:
'''
#(BGR)
# detected is green
for box in recognized_boxes:
cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 3)
# ground truth is blue
for box in gt_boxes:
cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (255, 0, 0), 3)
cv2.imwrite(outpath, image)
def save_boxes(args, recognized_boxes, recognized_scores, img_id):
if len(recognized_scores) < 1 and len(recognized_boxes) < 1:
return
pdf_name = img_id.split("/")[0]
math_csv_path = os.path.join(args.save_folder, args.exp_name, pdf_name + ".csv")
if not os.path.exists(os.path.dirname(math_csv_path)):
os.makedirs(os.path.dirname(math_csv_path))
math_output = open(math_csv_path, 'a')
recognized_boxes = np.concatenate((recognized_boxes,np.transpose([recognized_scores])),axis=1)
page_num = int(img_id.split("/")[-1])
col = np.array([int(page_num) - 1] * recognized_boxes.shape[0])
math_regions = np.concatenate((col[:, np.newaxis], recognized_boxes), axis=1)
np.savetxt(math_output, math_regions, fmt='%.2f', delimiter=',')
math_output.close()
#
#
# for i, box in enumerate(recognized_boxes):
# math_output.write(str(box[0]) + ',' + str(box[1]) + ',' + str(box[2]) + ',' +
# str(box[3]) + ',' + str(recognized_scores[i]) + '\n')
#
def draw_boxes(args, im, recognized_boxes, recognized_scores, boxes, confs, scale, img_id):
path = os.path.join("eval", args.exp_name, img_id + ".png")
if not os.path.exists(os.path.dirname(path)):
os.makedirs(os.path.dirname(path))
# Create figure and axes
fig,ax = plt.subplots(1)
scale = scale.cpu().numpy()
# Display the image
ax.imshow(im)
width, height, channels = im.shape
heatmap = np.zeros([width, height])
if len(recognized_scores) > 1 and len(recognized_boxes) > 1:
# Recognition heatmap
data = np.concatenate((recognized_boxes,np.transpose([recognized_scores])),axis=1)
data = data[data[:, 4].argsort()]
for box in data:
heatmap[int(box[1]):int(box[3]), int(box[0]):int(box[2])] = box[4]
for box in recognized_boxes:
rect = patches.Rectangle((box[0], box[1]), box[2]-box[0], box[3] - box[1],
linewidth=1, edgecolor='g', facecolor='none')
#Add the patch to the Axes
ax.add_patch(rect)
# Following line makes sure that all the heatmaps are in the scale, 0 to 1
# So color assigned to different scores are consistent across heatmaps for
# different images
heatmap[0:1, 0:1] = 1
heatmap[0:1, 1:2] = 0
plt.imshow(heatmap, alpha=0.4, cmap='hot', interpolation='nearest')
plt.colorbar()
plt.title(args.exp_name)
plt.show()
plt.savefig(path, dpi=600)
plt.close()
if __name__ == "__main__":
draw_boxes() |