hugmodelapi / main.py
Alex Vega
up
e373c5a
Raw
History Blame Contribute Delete
2.24 kB
from fastapi import FastAPI, File, UploadFile
from transformers import AutoImageProcessor, SiglipForImageClassification
from PIL import Image
import torch
import io
MODEL_NAME = "prithivMLmods/Alphabet-Sign-Language-Detection"
try:
print(f"Cargando modelo '{MODEL_NAME}'...")
processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
model = SiglipForImageClassification.from_pretrained(MODEL_NAME)
print(f"Modelo '{MODEL_NAME}' cargado exitosamente")
model_loaded = True
except Exception as e:
print(f"Error al cargar el modelo: {e}")
print("Usando modelo de ejemplo. Para uso real, necesitas un modelo específico de ASL.")
model = None
processor = None
model_loaded = False
app = FastAPI(title="API de ASL con modelo de HF")
@app.post("/predict/")
async def translate_sign(file: UploadFile = File(...)):
if not model or not processor:
return {
"error": "Modelo no disponible.",
"message": "El modelo específico de ASL no pudo cargarse. Verifica que el modelo existe en Hugging Face.",
"model_attempted": MODEL_NAME
}
try:
image_bytes = await file.read()
image = Image.open(io.BytesIO(image_bytes))
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
probs = torch.nn.functional.softmax(logits, dim=1)
predicted_class_idx = logits.argmax(-1).item()
confidence = probs[0][predicted_class_idx].item()
predicted_label = model.config.id2label[predicted_class_idx]
return {
"prediction": predicted_label,
"confidence": round(confidence * 100, 2)
}
except Exception as e:
return {"error": f"Error al procesar la imagen: {str(e)}"}
@app.get("/")
def read_root():
return {"message": "API ok. Usa el endpoint /predict/ para predecir."}
@app.get("/status/")
def get_status():
return {
"model_loaded": model_loaded,
"model_name": MODEL_NAME,
"message": "Modelo cargado correctamente" if model_loaded else "Modelo no disponible"
}