Spaces:
Running
Running
| from transformers import AutoImageProcessor, SuperPointForKeypointDetection | |
| import torch | |
| import matplotlib.pyplot as plt | |
| import uuid | |
| import gradio as gr | |
| processor = AutoImageProcessor.from_pretrained("magic-leap-community/superpoint") | |
| model = SuperPointForKeypointDetection.from_pretrained("magic-leap-community/superpoint") | |
| def infer(image): | |
| inputs = processor(image,return_tensors="pt").to(model.device, model.dtype) | |
| model_outputs = model(**inputs) | |
| image_sizes = [(image.size[1], image.size[0])] | |
| outputs = processor.post_process_keypoint_detection(model_outputs, image_sizes) | |
| keypoints = outputs[0]["keypoints"].detach().numpy() | |
| scores = outputs[0]["scores"].detach().numpy() | |
| image_width, image_height = image.size | |
| plt.axis('off') | |
| plt.imshow(image) | |
| plt.scatter( | |
| keypoints[:, 0], | |
| keypoints[:, 1], | |
| s=scores * 100, | |
| c='cyan', | |
| alpha=0.4 | |
| ) | |
| path = "./" + uuid.uuid4().hex + ".png" | |
| plt.savefig(path) | |
| plt.close() | |
| return path | |
| title = "SuperPoint" | |
| description = "Try [SuperPoint](https://huggingface.co/docs/transformers/en/model_doc/superpoint) in this demo, foundation model for keypoint detection supported in 🤗 transformers. Simply upload an image or try the example. " | |
| iface = gr.Interface(fn = infer, inputs = gr.Image(type="pil"), | |
| outputs = gr.Image(), title=title, description=description, examples=["./bee.jpg"]) | |
| iface.launch() |