yonadab commited on
Commit
478805f
verified
1 Parent(s): f6de5bc

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +40 -0
app.py ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, File, UploadFile
2
+ from fastapi.middleware.cors import CORSMiddleware
3
+ from transformers import AutoImageProcessor, AutoModel
4
+ from PIL import Image
5
+ import torch
6
+ import io
7
+
8
+ app = FastAPI()
9
+
10
+ # CORS (para pruebas locales o producci贸n cruzada)
11
+ app.add_middleware(
12
+ CORSMiddleware,
13
+ allow_origins=["*"], # Cambia esto en producci贸n
14
+ allow_credentials=True,
15
+ allow_methods=["*"],
16
+ allow_headers=["*"],
17
+ )
18
+
19
+ # Carga del modelo y procesador
20
+ processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
21
+ model = AutoModel.from_pretrained("facebook/dinov2-base")
22
+ model.eval()
23
+
24
+ @app.post("/embedding")
25
+ async def get_embedding(file: UploadFile = File(...)):
26
+ try:
27
+ image_bytes = await file.read()
28
+ image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
29
+
30
+ inputs = processor(images=image, return_tensors="pt")
31
+ with torch.no_grad():
32
+ outputs = model(**inputs)
33
+
34
+ # Promedio de los embeddings de todos los tokens (sin CLS)
35
+ embedding = outputs.last_hidden_state.mean(dim=1).squeeze().tolist()
36
+
37
+ return {"embedding": embedding}
38
+
39
+ except Exception as e:
40
+ return {"error": str(e)}