app.py

import torch
from torch import nn
from transformers import AutoImageProcessor, AutoModel
from PIL import Image
import base64
from io import BytesIO
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

# Nome do modelo no Hugging Face Hub
MODEL_NAME = "facebook/dinov2-small"

# Carregando processador e modelo
processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME)

# Projeção para 512D
projection = nn.Linear(model.config.hidden_size, 512)

# Inicializa o FastAPI
app = FastAPI(
    title="API de Embedding de Imagem",
    description="Endpoint para obter o embedding de uma imagem usando o modelo DINOv2.",
    version="1.0.0"
)

# Define o modelo de dados para a requisição
class ImageRequest(BaseModel):
    image: str

# Define o endpoint para o embedding da imagem
@app.post("/embed")
async def get_embedding(request: ImageRequest):
    try:
        # Extrai a string Base64 do formato "data:image/png;base64,..."
        header, img_base64 = request.image.split(",", 1)
        
        # Decodifica a string Base64
        image_data = base64.b64decode(img_base64)
        
        # Abre a imagem com Pillow
        image = Image.open(BytesIO(image_data))
        
        # Preprocessamento
        inputs = processor(images=image, return_tensors="pt")
        
        with torch.no_grad():
            outputs = model(**inputs)
            last_hidden_state = outputs.last_hidden_state
            embedding = last_hidden_state[:, 0]
            embedding_512 = projection(embedding)
            
        # Converte para lista Python e retorna
        return {"embedding": embedding_512.squeeze().tolist()}
        
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Erro ao processar a imagem: {e}")