aresca commited on
Commit
b9ea15d
verified
1 Parent(s): 6c73872

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +28 -0
app.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, UploadFile, File
2
+ from transformers import AutoModelForImageClassification, AutoProcessor
3
+ from PIL import Image
4
+ import torch
5
+
6
+ app = FastAPI()
7
+
8
+ # Carga el modelo y el preprocesador desde Hugging Face
9
+ model_name = "jazzmacedo/fruits-and-vegetables-detector-36"
10
+ model = AutoModelForImageClassification.from_pretrained(model_name)
11
+ processor = AutoProcessor.from_pretrained(model_name)
12
+
13
+ @app.post("/predict")
14
+ async def predict(file: UploadFile = File(...)):
15
+ # Procesa la imagen recibida
16
+ image = Image.open(file.file).convert("RGB")
17
+ inputs = processor(images=image, return_tensors="pt")
18
+
19
+ # Realiza la predicci贸n
20
+ with torch.no_grad():
21
+ outputs = model(**inputs)
22
+
23
+ # Obtiene la predicci贸n con mayor probabilidad
24
+ logits = outputs.logits
25
+ predicted_class_idx = logits.argmax(-1).item()
26
+ predicted_class_name = model.config.id2label[predicted_class_idx]
27
+
28
+ return {"prediction": predicted_class_name}