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() |