Fix keypoint visualization code

README.md

@@ -78,23 +78,47 @@ model = SuperPointForKeypointDetection.from_pretrained("magic-leap-community/sup
 
 inputs = processor(images, return_tensors="pt")
 outputs = model(**inputs)
+```
+
+We can now visualize the keypoints.
+
+```
+import matplotlib.pyplot as plt
+import torch
 
 for i in range(len(images)):
+    image = images[i]
+    image_width, image_height = image.size  
+
     image_mask = outputs.mask[i]
     image_indices = torch.nonzero(image_mask).squeeze()
-    image_keypoints = outputs.keypoints[i][image_indices]
-    image_scores = outputs.scores[i][image_indices]
-    image_descriptors = outputs.descriptors[i][image_indices]
-```
 
-You can then print the keypoints on the image to visualize the result :
-```python
-import cv2
-for keypoint, score in zip(image_keypoints, image_scores):
-    keypoint_x, keypoint_y = int(keypoint[0].item()), int(keypoint[1].item())
-    color = tuple([score.item() * 255] * 3)
-    image = cv2.circle(image, (keypoint_x, keypoint_y), 2, color)
-cv2.imwrite("output_image.png", image)
+    image_scores = outputs.scores[i][image_indices]
+    image_keypoints = outputs.keypoints[i][image_indices]
+    
+    keypoints = image_keypoints.detach().numpy()
+    scores = image_scores.detach().numpy()
+
+    valid_keypoints = [
+        (kp, score) for kp, score in zip(keypoints, scores)
+        if 0 <= kp[0] < image_width and 0 <= kp[1] < image_height
+    ]
+
+    valid_keypoints, valid_scores = zip(*valid_keypoints)
+    valid_keypoints = torch.tensor(valid_keypoints)
+    valid_scores = torch.tensor(valid_scores)
+
+    print(valid_keypoints.shape)
+
+    plt.axis('off')
+    plt.imshow(image)
+    plt.scatter(
+        valid_keypoints[:, 0], 
+        valid_keypoints[:, 1], 
+        s=valid_scores * 100, 
+        c='red'
+    )
+    plt.show()
 ```
 
 This model was contributed by [stevenbucaille](https://huggingface.co/stevenbucaille).