Spaces:
Sleeping
Sleeping
File size: 2,763 Bytes
948f9ef | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | import io
import numpy as np
import tensorflow as tf
from fastapi import FastAPI, File, UploadFile, HTTPException
from PIL import Image
from typing import Dict
app = FastAPI(title="DenseNet121 Image Prediction API")
# Labels dictionary as per requirements
LABELS = {
0: {"short": "DR", "full": "Diabetic Retinopathy"},
1: {"short": "MH", "full": "Media Haze"},
2: {"short": "NL", "full": "Normal"}
}
# Global variable to store the model
model = None
@app.on_event("startup")
async def load_model():
global model
try:
# Loading the fully trained model
model = tf.keras.models.load_model("DenseNet121.h5")
print("Model loaded successfully.")
except Exception as e:
print(f"Error loading model: {e}")
def preprocess_image(image: Image.Image) -> np.ndarray:
"""
Preprocess the image according to requirements:
- RGB format
- Resize to 224x224
- Normalize by dividing by 255
"""
if image.mode != "RGB":
image = image.convert("RGB")
image = image.resize((224, 224))
img_array = np.array(image)
# Normalize by 255
img_array = img_array.astype(np.float32) / 255.0
# Add batch dimension
img_array = np.expand_dims(img_array, axis=0)
return img_array
@app.get("/")
async def root():
return {"message": "DenseNet121 Prediction API is running", "model_status": "Loaded" if model else "Not Loaded"}
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
"""
Endpoint for image prediction.
Accepts an image file and returns softmax probabilities and class labels.
"""
if model is None:
raise HTTPException(status_code=500, detail="Model not loaded.")
try:
# Read image
contents = await file.read()
image = Image.open(io.BytesIO(contents))
# Preprocess
processed_image = preprocess_image(image)
# Predict
predictions = model.predict(processed_image)[0] # Softmax probabilities
# Prepare response
results = {}
summary_parts = []
for i, prob in enumerate(predictions):
label_info = LABELS[i]
percentage = prob * 100
results[label_info["full"]] = f"{percentage:.2f}%"
summary_parts.append(f"{int(round(percentage))}% {label_info['short']}")
summary_text = ", ".join(summary_parts)
return {
"prediction_summary": summary_text,
"detailed_probabilities": results,
"top_prediction": LABELS[np.argmax(predictions)]["full"]
}
except Exception as e:
raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")
|