Spaces:
Running
Running
from fastapi import FastAPI, UploadFile, File | |
from fastapi.middleware.cors import CORSMiddleware | |
import numpy as np | |
from tensorflow.keras.models import load_model | |
from tensorflow.keras.preprocessing.image import load_img, img_to_array | |
from io import BytesIO | |
app = FastAPI() | |
# Enable CORS to allow requests from frontend (React) | |
app.add_middleware( | |
CORSMiddleware, | |
allow_origins=["*"], # Change ["http://localhost:5173"] for better security | |
allow_credentials=True, | |
allow_methods=["*"], | |
allow_headers=["*"], | |
) | |
# Load your model | |
model = load_model("densenet201_food_classification.h5") | |
# Define class indices | |
class_indices = { | |
0: "burger", | |
1: "butter_naan", | |
2: "chai", | |
3: "chapati", | |
4: "chole_bhature", | |
5: "dal_makhani", | |
6: "dhokla", | |
7: "fried_rice", | |
8: "idli", | |
9: "jalebi", | |
10: "kaathi_rolls", | |
11: "kadai_paneer", | |
12: "kulfi", | |
13: "masala_dosa", | |
14: "momos", | |
15: "paani_puri", | |
16: "pakode", | |
17: "pav_bhaji", | |
18: "pizza", | |
19: "samosa" | |
} | |
def predict_image(image, model): | |
try: | |
img = load_img(image, target_size=(224, 224)) | |
image_array = img_to_array(img) / 255.0 | |
image_array = np.expand_dims(image_array, axis=0) | |
predictions = model.predict(image_array) | |
class_idx = np.argmax(predictions) | |
class_label = class_indices.get(class_idx, "Unknown") | |
confidence = float(predictions[0][class_idx]) | |
return class_label, confidence | |
except Exception as e: | |
return None, None | |
async def predict(file: UploadFile = File(...)): | |
try: | |
image_data = await file.read() | |
image = BytesIO(image_data) | |
class_label, confidence = predict_image(image, model) | |
if class_label is None: | |
return {"error": "Prediction failed"} | |
return {"predicted_class": class_label, "confidence": f"{confidence:.2f}"} | |
except Exception as e: | |
return {"error": f"Internal Server Error: {str(e)}"} | |