Alex Vega commited on
Commit
e373c5a
·
1 Parent(s): 212f1ad
Files changed (1) hide show
  1. main.py +43 -17
main.py CHANGED
@@ -1,23 +1,27 @@
1
  from fastapi import FastAPI, File, UploadFile
2
- from transformers import ViTImageProcessor, ViTForImageClassification
3
  from PIL import Image
4
  import torch
5
  import io
6
 
7
- MODEL_NAME = "ahmed-masoud/sign_language_translator"
8
 
9
 
10
  try:
11
- processor = ViTImageProcessor.from_pretrained(MODEL_NAME)
 
12
 
13
- model = ViTForImageClassification.from_pretrained(MODEL_NAME)
14
 
15
- print(f"Modelo '{MODEL_NAME}' cargado")
 
16
 
17
  except Exception as e:
18
- print(f"Error al cargar el modelo {e}")
 
19
  model = None
20
  processor = None
 
21
 
22
  app = FastAPI(title="API de ASL con modelo de HF")
23
 
@@ -25,24 +29,46 @@ app = FastAPI(title="API de ASL con modelo de HF")
25
  @app.post("/predict/")
26
  async def translate_sign(file: UploadFile = File(...)):
27
  if not model or not processor:
28
- return {"error": "Modelo no disponible."}
 
 
 
 
29
 
30
- image_bytes = await file.read()
31
- image = Image.open(io.BytesIO(image_bytes))
 
32
 
33
- inputs = processor(images=image, return_tensors="pt")
34
 
35
- with torch.no_grad():
36
- outputs = model(**inputs)
37
- logits = outputs.logits
 
 
38
 
39
- predicted_class_idx = logits.argmax(-1).item()
40
-
41
- predicted_label = model.config.id2label[predicted_class_idx]
 
42
 
43
- return {"prediction": predicted_label}
 
 
 
 
 
 
44
 
45
 
46
  @app.get("/")
47
  def read_root():
48
  return {"message": "API ok. Usa el endpoint /predict/ para predecir."}
 
 
 
 
 
 
 
 
 
1
  from fastapi import FastAPI, File, UploadFile
2
+ from transformers import AutoImageProcessor, SiglipForImageClassification
3
  from PIL import Image
4
  import torch
5
  import io
6
 
7
+ MODEL_NAME = "prithivMLmods/Alphabet-Sign-Language-Detection"
8
 
9
 
10
  try:
11
+ print(f"Cargando modelo '{MODEL_NAME}'...")
12
+ processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
13
 
14
+ model = SiglipForImageClassification.from_pretrained(MODEL_NAME)
15
 
16
+ print(f"Modelo '{MODEL_NAME}' cargado exitosamente")
17
+ model_loaded = True
18
 
19
  except Exception as e:
20
+ print(f"Error al cargar el modelo: {e}")
21
+ print("Usando modelo de ejemplo. Para uso real, necesitas un modelo específico de ASL.")
22
  model = None
23
  processor = None
24
+ model_loaded = False
25
 
26
  app = FastAPI(title="API de ASL con modelo de HF")
27
 
 
29
  @app.post("/predict/")
30
  async def translate_sign(file: UploadFile = File(...)):
31
  if not model or not processor:
32
+ return {
33
+ "error": "Modelo no disponible.",
34
+ "message": "El modelo específico de ASL no pudo cargarse. Verifica que el modelo existe en Hugging Face.",
35
+ "model_attempted": MODEL_NAME
36
+ }
37
 
38
+ try:
39
+ image_bytes = await file.read()
40
+ image = Image.open(io.BytesIO(image_bytes))
41
 
42
+ inputs = processor(images=image, return_tensors="pt")
43
 
44
+ with torch.no_grad():
45
+ outputs = model(**inputs)
46
+ logits = outputs.logits
47
+
48
+ probs = torch.nn.functional.softmax(logits, dim=1)
49
 
50
+ predicted_class_idx = logits.argmax(-1).item()
51
+ confidence = probs[0][predicted_class_idx].item()
52
+
53
+ predicted_label = model.config.id2label[predicted_class_idx]
54
 
55
+ return {
56
+ "prediction": predicted_label,
57
+ "confidence": round(confidence * 100, 2)
58
+ }
59
+
60
+ except Exception as e:
61
+ return {"error": f"Error al procesar la imagen: {str(e)}"}
62
 
63
 
64
  @app.get("/")
65
  def read_root():
66
  return {"message": "API ok. Usa el endpoint /predict/ para predecir."}
67
+
68
+ @app.get("/status/")
69
+ def get_status():
70
+ return {
71
+ "model_loaded": model_loaded,
72
+ "model_name": MODEL_NAME,
73
+ "message": "Modelo cargado correctamente" if model_loaded else "Modelo no disponible"
74
+ }