Update README.md

README.md

@@ -130,45 +130,59 @@ image0_matching_scores = outputs.matching_scores[0, 0][image0_indices]
 image1_matching_scores = outputs.matching_scores[0, 1][image1_indices]
 ```
 
-You can then print the matched keypoints on a side-by-side image to visualize the result :
+You can use the `post_process_keypoint_matching` method from the `SuperGlueImageProcessor` to get the keypoints and matches in a more readable format:
 ```python
-import cv2
+image_sizes = [(image.height, image.width) for image in images]
+outputs = processor.post_process_keypoint_matching(outputs, image_sizes)
+for i, output in enumerate(outputs):
+    print("For the image pair", i)
+    for keypoint0, keypoint1, matching_score in zip(output["keypoints0"], output["keypoints1"],
+                                                    output["matching_scores"]):
+        print(
+            f"Keypoint at coordinate {keypoint0.numpy()} in the first image matches with keypoint at coordinate {keypoint1.numpy()} in the second image with a score of {matching_score}."
+        )
+```
+
+From the outputs, you can visualize the matches between the two images using the following code:
+```python
+import matplotlib.pyplot as plt
 import numpy as np
 
 # Create side by side image
-input_data = inputs['pixel_values']
-height, width = input_data.shape[-2:]
-matched_image = np.zeros((height, width * 2, 3))
-matched_image[:, :width] = input_data.squeeze()[0].permute(1, 2, 0).cpu().numpy()
-matched_image[:, width:] = input_data.squeeze()[1].permute(1, 2, 0).cpu().numpy()
-matched_image = (matched_image * 255).astype(np.uint8)
-
-# Retrieve matches by looking at which keypoints in image0 actually matched with keypoints in image1
-image0_mask = outputs.mask[0, 0]
-image0_indices = torch.nonzero(image0_mask).squeeze()
-image0_matches_indices = torch.nonzero(outputs.matches[0, 0][image0_indices] != -1).squeeze()
-image0_keypoints = outputs.keypoints[0, 0][image0_matches_indices]
-image0_matches = outputs.matches[0, 0][image0_matches_indices]
-image0_matching_scores = outputs.matching_scores[0, 0][image0_matches_indices]
-# Retrieve matches from image1
-image1_mask = outputs.mask[0, 1]
-image1_indices = torch.nonzero(image1_mask).squeeze()
-image1_keypoints = outputs.keypoints[0, 1][image0_matches]
-
-# Draw matches
-for keypoint0, keypoint1, score in zip(image0_keypoints, image1_keypoints, image0_matching_scores):
-    keypoint0_x, keypoint0_y = int(keypoint0[0]), int(keypoint0[1])
-    keypoint1_x, keypoint1_y = int(keypoint1[0] + width), int(keypoint1[1])
-    color = [0, 1, 0, 0.5]  # Set color based on score
-    plt.plot([keypoint0_x, keypoint1_x], [keypoint0_y, keypoint1_y], color=color, linewidth=1)
-
-# Save the image
+merged_image = np.zeros((max(image1.height, image2.height), image1.width + image2.width, 3))
+merged_image[: image1.height, : image1.width] = np.array(image1) / 255.0
+merged_image[: image2.height, image1.width :] = np.array(image2) / 255.0
+plt.imshow(merged_image)
+plt.axis("off")
+
+# Retrieve the keypoints and matches
+output = outputs[0]
+keypoints0 = output["keypoints0"]
+keypoints1 = output["keypoints1"]
+matching_scores = output["matching_scores"]
+keypoints0_x, keypoints0_y = keypoints0[:, 0].numpy(), keypoints0[:, 1].numpy()
+keypoints1_x, keypoints1_y = keypoints1[:, 0].numpy(), keypoints1[:, 1].numpy()
+
+# Plot the matches
+for keypoint0_x, keypoint0_y, keypoint1_x, keypoint1_y, matching_score in zip(
+        keypoints0_x, keypoints0_y, keypoints1_x, keypoints1_y, matching_scores
+):
+    plt.plot(
+        [keypoint0_x, keypoint1_x + image1.width],
+        [keypoint0_y, keypoint1_y],
+        color=plt.get_cmap("RdYlGn")(matching_score.item()),
+        alpha=0.9,
+        linewidth=0.5,
+    )
+    plt.scatter(keypoint0_x, keypoint0_y, c="black", s=2)
+    plt.scatter(keypoint1_x + image1.width, keypoint1_y, c="black", s=2)
+
+# Save the plot
 plt.savefig("matched_image.png", dpi=300, bbox_inches='tight')
 plt.close()
 ```
 
-![image/png](https://cdn-uploads.huggingface.co/production/uploads/632885ba1558dac67c440aa8/PiLL7svnN2dTqOxrsobJb.png)
-
+![image/png](https://cdn-uploads.huggingface.co/production/uploads/632885ba1558dac67c440aa8/01ZYaLB1NL5XdA8u7yCo4.png)
 
 ## Training Details
 
@@ -193,7 +207,7 @@ The model has 12 million parameters, making it relatively compact compared to so
 The inference speed of SuperGlue is suitable for real-time applications and can be readily integrated into 
 modern Simultaneous Localization and Mapping (SLAM) or Structure-from-Motion (SfM) systems.
 
-## Citation [optional]
+## Citation
 
 <!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->