ropaembeddings / app.py
yonadab's picture
Update app.py
506a05e verified
raw
history blame
818 Bytes
import gradio as gr
from transformers import AutoImageProcessor, AutoModel
import torch
from PIL import Image
import base64
import io
# Modelo
processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
model = AutoModel.from_pretrained("facebook/dinov2-base")
# Recibe string base64
def get_embedding(base64_str):
# Separar encabezado
header, encoded = base64_str.split(",", 1)
# Decodificar
image_bytes = base64.b64decode(encoded)
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
# Procesar
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
embeddings = model(**inputs).last_hidden_state[:, 0]
return embeddings.squeeze().tolist()
# Gradio Interface
gr.Interface(fn=get_embedding, inputs="text", outputs="json").launch()