Update app.py

app.py

@@ -1,11 +1,9 @@
 import torch
 from torch import nn
 from transformers import AutoImageProcessor, AutoModel
+import gradio as gr
+import numpy as np
 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"
@@ -14,47 +12,33 @@ MODEL_NAME = "facebook/dinov2-small"
 processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
 model = AutoModel.from_pretrained(MODEL_NAME)
 
-# Projeção para 512D
+# Projeção para 512D (caso a saída seja >512, reduzimos)
 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"
+def get_embedding(image: Image.Image):
+    # Preprocessamento
+    inputs = processor(images=image, return_tensors="pt")
+    
+    with torch.no_grad():
+        outputs = model(**inputs)
+        # Usando o CLS token como embedding da imagem
+        last_hidden_state = outputs.last_hidden_state  # (batch, seq_len, hidden)
+        embedding = last_hidden_state[:, 0]  # pegando o [CLS] token
+
+        # Projeta para 512D
+        embedding_512 = projection(embedding)
+
+    # Converte para lista Python
+    return embedding_512.squeeze().tolist()
+
+# Cria API com Gradio (sem interface visual, apenas endpoint)
+iface = gr.Interface(
+    fn=get_embedding,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.JSON(),
+    live=False,
+    api_name="embed"  # endpoint em /embed
 )
 
-# Define o modelo de dados para a requisição
-class ImageRequest(BaseModel):
-    image: str  # A imagem é uma string em formato "data:image/png;base64,..."
-
-# 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}")
-
-# Para executar, salve o arquivo como `main.py` e use o comando:
-# uvicorn main:app --reload
+if __name__ == "__main__":
+    iface.launch()