mirrors/mmocr
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmocr.git synced 2025-06-03 21:54:47 +08:00
Code Issues Projects Releases Wiki Activity

[Bug] Bug generated during kie inference visualization (#1830)

Browse Source 
* Update kie_visualizer.py

* Update kie_visualizer.py

* Update kie_visualizer.py
This commit is contained in:
YangLy 2023-04-03 17:39:02 +08:00 committed by GitHub
parent 4842599191
commit e6174b29fe
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mmocr/visualization/kie_visualizer.py
View File

@ -70,6 +70,8 @@ class KIELocalVisualizer(BaseLocalVisualizer):
np.ndarray: The image with edge labels drawn.
"""
pairs = np.where(edge_labels > 0)
if torch.is_tensor(pairs):
pairs = pairs.cpu()
key_bboxes = bboxes[pairs[0]]
value_bboxes = bboxes[pairs[1]]
x_data = np.stack([(key_bboxes[:, 2] + key_bboxes[:, 0]) / 2,