File size: 1,953 Bytes
c82a3d1 478805f c82a3d1 478805f c82a3d1 478805f c82a3d1 478805f c82a3d1 533df08 478805f c82a3d1 478805f c82a3d1 478805f c82a3d1 478805f c82a3d1 478805f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from transformers import AutoImageProcessor, AutoModel
from PIL import Image
import torch
import uuid
import io
app = FastAPI()
# Habilita CORS si lo necesitas
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Carga del modelo
processor = AutoImageProcessor.from_pretrained("facebook/dinov2-small")
model = AutoModel.from_pretrained("facebook/dinov2-small")
model.eval()
# Memoria temporal para almacenar imágenes (podrías usar base de datos si prefieres)
temp_images = {}
event_ids = {}
# Paso 1: Subida de imagen + event_id
@app.post("/upload")
async def upload_image(file: UploadFile = File(...), event_id: str = Form(...)):
try:
content = await file.read()
image_id = str(uuid.uuid4())
temp_images[image_id] = content
event_ids[image_id] = event_id
return {"image_id": image_id, "event_id": event_id}
except Exception as e:
return {"error": str(e)}
# Paso 2: Obtener embedding por image_id
@app.post("/embedding")
async def get_embedding(image_id: str = Form(...)):
if image_id not in temp_images:
raise HTTPException(status_code=404, detail="image_id not found")
event_id = event_ids[image_id]
image_bytes = temp_images[image_id]
try:
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
# Promedio de todos los tokens (puedes cambiar por CLS si quieres)
embedding = outputs.last_hidden_state.mean(dim=1).squeeze().tolist()
return {
"event_id": event_id,
"embedding": embedding
}
except Exception as e:
return {"error": str(e)}
|