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

* Update kie_visualizer.py * Update kie_visualizer.py * Update kie_visualizer.py
2023-04-03 17:39:02 +08:00 · 2023-04-03 17:39:02 +08:00 · e6174b29fe
parent 4842599191
commit e6174b29fe
1 changed files with 2 additions and 0 deletions
--- a/mmocr/visualization/kie_visualizer.py
+++ b/mmocr/visualization/kie_visualizer.py
@ -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,