Update app.py

app.py

@@ -6,47 +6,38 @@ import base64
 from io import BytesIO
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
+import imagehash  # <-- NOVO
 
-# Nome do modelo no Hugging Face Hub
 MODEL_NAME = "facebook/dinov2-small"
-
-# Carregando processador e modelo
-# Usamos a classe específica Dinov2Model para garantir que o modelo seja carregado corretamente
 processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
 model = Dinov2Model.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.",
+    description="Endpoint para obter o embedding e pHash de uma imagem.",
     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:
         header, img_base64 = request.image.split(",", 1)
         image_data = base64.b64decode(img_base64)
-        image = Image.open(BytesIO(image_data))
-        
-        # --- Lógica de Inferência do seu script original ---
+        image = Image.open(BytesIO(image_data)).convert("RGB")
         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)
-            
-        return {"embedding": embedding_512.squeeze().tolist()}
-        
+        phash = str(imagehash.phash(image))  # <-- NOVO
+        return {
+            "embedding": embedding_512.squeeze().tolist(),
+            "phash": phash
+        }
     except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Erro ao processar a imagem: {e}")
+        raise HTTPException(status_code=400, detail=f"Erro ao processar a imagem: {e}")