File size: 3,726 Bytes
f361235 | 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 96 97 98 99 | from fastapi import FastAPI, HTTPException, Depends, status
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from pydantic import BaseModel, Field
import joblib
import os
import numpy as np
# Initialize FastAPI app
app = FastAPI(
title="Iris Classification API",
description="A REST API for predicting Iris species using a pre-trained scikit-learn model.",
version="1.0.0"
)
# --- Authentication Setup ---
security = HTTPBasic()
def get_current_username(credentials: HTTPBasicCredentials = Depends(security)):
correct_username = os.getenv("API_USERNAME")
correct_password = os.getenv("API_PASSWORD")
if not correct_username or not correct_password:
# This handles cases where secrets aren't set in HF Spaces (shouldn't happen if done correctly)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="API credentials not configured on the server."
)
if not (credentials.username == correct_username and credentials.password == correct_password):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Incorrect username or password",
headers={"WWW-Authenticate": "Basic"},
)
return credentials.username
# --- Model Loading ---
model = None
class_names = None
@app.on_event("startup")
async def load_artifacts():
global model, class_names
model_path = os.path.join("model", "iris_model.joblib")
class_names_path = os.path.join("model", "iris_class_names.joblib")
if not os.path.exists(model_path) or not os.path.exists(class_names_path):
raise RuntimeError(f"Model or class names file not found. Ensure '{model_path}' and '{class_names_path}' exist.")
model = joblib.load(model_path)
class_names = joblib.load(class_names_path)
print("Model and class names loaded successfully.")
# --- Request Body Model ---
class IrisFeatures(BaseModel):
sepal_length: float = Field(..., example=5.1, description="Sepal length in cm")
sepal_width: float = Field(..., example=3.5, description="Sepal width in cm")
petal_length: float = Field(..., example=1.4, description="Petal length in cm")
petal_width: float = Field(..., example=0.2, description="Petal width in cm")
# --- API Endpoint ---
@app.post("/predict", summary="Predict Iris Species", response_description="The predicted Iris species and probabilities.")
async def predict_iris(
features: IrisFeatures,
current_user: str = Depends(get_current_username)
):
if model is None or class_names is None:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Model is not loaded yet. Please try again in a moment."
)
input_data = np.array([[
features.sepal_length,
features.sepal_width,
features.petal_length,
features.petal_width
]])
prediction_index = model.predict(input_data)[0]
predicted_species = class_names[prediction_index]
probabilities = model.predict_proba(input_data)[0]
probabilities_dict = {name: float(prob) for name, prob in zip(class_names, probabilities)}
return {
"predicted_species": predicted_species,
"prediction_probabilities": probabilities_dict
}
# --- Health Check Endpoint ---
@app.get("/health", summary="Health Check", response_description="Indicates if the API is running.")
async def health_check():
return {"status": "ok", "model_loaded": model is not None}
# Note: The uvicorn.run part is for local execution.
# Hugging Face Spaces will use the CMD in the Dockerfile.
# For local testing in Colab, you'd use ngrok or colabcode (see below).
|