ropaembeddings / app.py
yonadab's picture
Create app.py
61bc9be verified
raw
history blame
739 Bytes
import gradio as gr
from transformers import AutoImageProcessor, AutoModel
import torch
# Cargar el modelo DINOv2 una sola vez
processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
model = AutoModel.from_pretrained("facebook/dinov2-base")
def get_embedding(image):
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
embeddings = model(**inputs).last_hidden_state[:, 0] # CLS token
return embeddings.squeeze().tolist()
# Gradio UI solo para aceptar imágenes y devolver JSON
iface = gr.Interface(
fn=get_embedding,
inputs=gr.Image(type="pil"),
outputs="json",
description="Microservicio para extraer embeddings de imágenes usando DINOv2."
)
iface.launch()