Update app.py

app.py

@@ -7,49 +7,75 @@ import joblib
 import pandas as pd
 from huggingface_hub import hf_hub_download
 import os
+import logging
 
+# ------------------------------------------------------------
+# Setup Logging for Debugging
+# ------------------------------------------------------------
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# ------------------------------------------------------------
+# FastAPI App Initialization
+# ------------------------------------------------------------
 app = FastAPI(title="PulmoProbe AI API")
 
-# Add CORS middleware to allow your frontend app to call the API
+# Enable CORS so React frontend can communicate with FastAPI backend
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=["*"],  # In production, specify domain
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# --- Set Hugging Face cache to a writable location ---
+# ------------------------------------------------------------
+# Hugging Face Model Setup
+# ------------------------------------------------------------
+# Use a writable temp directory for Hugging Face cache
 os.environ['HF_HOME'] = '/tmp/huggingface'
 os.makedirs(os.environ['HF_HOME'], exist_ok=True)
-print(f"HF_HOME set to {os.environ['HF_HOME']}")
+logger.info(f"HF_HOME set to {os.environ['HF_HOME']}")
 
-# --- Download and Load Model from Hugging Face Hub ---
 MODEL_REPO_ID = "costaspinto/PulmoProbe"
 MODEL_FILENAME = "best_model.joblib"
 
-print("Downloading model from Hugging Face Hub...")
-model_path = hf_hub_download(
-    repo_id=MODEL_REPO_ID,
-    filename=MODEL_FILENAME,
-    cache_dir=os.environ['HF_HOME']
-)
-model = joblib.load(model_path)
-print("Model loaded successfully.")
+logger.info("Downloading model from Hugging Face Hub...")
+try:
+    model_path = hf_hub_download(
+        repo_id=MODEL_REPO_ID,
+        filename=MODEL_FILENAME,
+        cache_dir=os.environ['HF_HOME']
+    )
+    model = joblib.load(model_path)
+    logger.info("Model loaded successfully.")
+except Exception as e:
+    logger.error(f"Failed to download or load model: {str(e)}")
+    raise RuntimeError(f"Model loading failed: {str(e)}")
 
-# --- Define Input Data Model for one-hot encoded features ---
-# This new model directly matches the one-hot encoded data from the frontend
+# ------------------------------------------------------------
+# Define Input Schema - Must Match Frontend Fields Exactly
+# ------------------------------------------------------------
 class OneHotPatientData(BaseModel):
+    # Continuous fields
     age: float
     bmi: float
     cholesterol_level: float
+
+    # Binary medical conditions
     hypertension: int
     asthma: int
     cirrhosis: int
     other_cancer: int
+
+    # Family history
     family_history_Yes: int
+
+    # Gender one-hot
     gender_Male: int
     gender_Female: int
+
+    # Country one-hot
     country_Sweden: int
     country_Netherlands: int
     country_Hungary: int
@@ -64,35 +90,61 @@ class OneHotPatientData(BaseModel):
     country_Spain: int
     country_UnitedKingdom: int
     country_UnitedStates: int
+
+    # Cancer stage one-hot
     cancer_stage_StageI: int
     cancer_stage_StageII: int
     cancer_stage_StageIII: int
     cancer_stage_StageIV: int
+
+    # Smoking status one-hot
     smoking_status_NeverSmoked: int
     smoking_status_FormerSmoker: int
     smoking_status_PassiveSmoker: int
     smoking_status_CurrentSmoker: int
+
+    # Treatment type one-hot
     treatment_type_Chemotherapy: int
     treatment_type_Surgery: int
     treatment_type_Radiation: int
     treatment_type_Combined: int
 
-# --- API Endpoints ---
+# ------------------------------------------------------------
+# Root Endpoint
+# ------------------------------------------------------------
 @app.get("/")
 def read_root():
     return {"message": "Welcome to the PulmoProbe AI API"}
 
+# ------------------------------------------------------------
+# Prediction Endpoint
+# ------------------------------------------------------------
 @app.post("/predict")
 def predict(data: OneHotPatientData):
     try:
-        input_df = pd.DataFrame([data.dict()])
+        # Convert incoming data to DataFrame
+        input_data = data.dict()
+        logger.info(f"Received prediction request: {input_data}")
+
+        input_df = pd.DataFrame([input_data])
+
+        # Make prediction
         probabilities = model.predict_proba(input_df)[0]
-        confidence_high_risk = probabilities[0]
+        confidence_high_risk = probabilities[0]  # Assuming index 0 = High Risk
+        logger.info(f"Model raw probabilities: {probabilities}")
+
+        # Determine risk level
         risk_level = "High Risk of Non-Survival" if confidence_high_risk > 0.5 else "Low Risk of Non-Survival"
 
-        return {
+        result = {
             "risk": risk_level,
             "confidence": f"{confidence_high_risk * 100:.1f}%"
         }
+
+        logger.info(f"Prediction result: {result}")
+        return result
+
     except Exception as e:
-        return {"error": str(e)}
+        logger.error(f"Prediction error: {str(e)}")
+        return {"error": str(e)}
+