Update app.py

app.py

@@ -10,20 +10,20 @@ import os
 import logging
 
 # ------------------------------------------------------------
-# Setup Logging for Debugging
+# Setup Logging
 # ------------------------------------------------------------
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # ------------------------------------------------------------
-# FastAPI App Initialization
+# FastAPI Initialization
 # ------------------------------------------------------------
 app = FastAPI(title="PulmoProbe AI API")
 
-# Enable CORS so React frontend can communicate with FastAPI backend
+# Allow CORS for frontend communication
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # In production, specify domain
+    allow_origins=["*"],  # Use specific domain in production
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -32,7 +32,6 @@ app.add_middleware(
 # ------------------------------------------------------------
 # 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)
 logger.info(f"HF_HOME set to {os.environ['HF_HOME']}")
@@ -50,11 +49,12 @@ try:
     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)}")
+    logger.error(f"Failed to load model: {str(e)}")
     raise RuntimeError(f"Model loading failed: {str(e)}")
 
 # ------------------------------------------------------------
-# Define Input Schema - Must Match Frontend Fields Exactly
+# Define Input Schema
+# Must exactly match the model feature names and order
 # ------------------------------------------------------------
 class OneHotPatientData(BaseModel):
     # Continuous fields
@@ -68,46 +68,54 @@ class OneHotPatientData(BaseModel):
     cirrhosis: int
     other_cancer: int
 
-    # Family history
-    family_history_Yes: int
-
-    # Gender one-hot
+    # Gender (Male = 1, Female = 0)
     gender_Male: int
-    gender_Female: int
 
-    # Country one-hot
-    country_Sweden: int
-    country_Netherlands: int
-    country_Hungary: int
+    # Countries (One-Hot)
     country_Belgium: int
-    country_Italy: int
+    country_Bulgaria: int
     country_Croatia: int
+    country_Cyprus: int
+    country_Czech_Republic: int
     country_Denmark: int
-    country_Germany: int
+    country_Estonia: int
+    country_Finland: int
     country_France: int
+    country_Germany: int
+    country_Greece: int
+    country_Hungary: int
+    country_Ireland: int
+    country_Italy: int
+    country_Latvia: int
+    country_Lithuania: int
+    country_Luxembourg: int
+    country_Malta: int
+    country_Netherlands: int
+    country_Poland: int
+    country_Portugal: int
+    country_Romania: int
     country_Slovakia: int
-    country_Finland: int
+    country_Slovenia: int
     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
+    country_Sweden: int
+
+    # Cancer stages (Stage I is baseline)
+    cancer_stage_Stage_Ii: int
+    cancer_stage_Stage_Iii: int
+    cancer_stage_Stage_Iv: int
+
+    # Family history
+    family_history_Yes: int
+
+    # Smoking status (Current Smoker is baseline)
+    smoking_status_Former_Smoker: int
+    smoking_status_Never_Smoked: int
+    smoking_status_Passive_Smoker: int
+
+    # Treatment type (Chemotherapy is baseline)
     treatment_type_Combined: int
+    treatment_type_Radiation: int
+    treatment_type_Surgery: int
 
 # ------------------------------------------------------------
 # Root Endpoint
@@ -122,18 +130,18 @@ def read_root():
 @app.post("/predict")
 def predict(data: OneHotPatientData):
     try:
-        # Convert incoming data to DataFrame
-        input_data = data.dict()
-        logger.info(f"Received prediction request: {input_data}")
+        input_dict = data.dict()
+        logger.info(f"Incoming data: {input_dict}")
 
-        input_df = pd.DataFrame([input_data])
+        # Create DataFrame
+        input_df = pd.DataFrame([input_dict])
 
-        # Make prediction
+        # Predict probabilities
         probabilities = model.predict_proba(input_df)[0]
-        confidence_high_risk = probabilities[0]  # Assuming index 0 = High Risk
-        logger.info(f"Model raw probabilities: {probabilities}")
+        logger.info(f"Model probabilities: {probabilities}")
 
-        # Determine risk level
+        # Assuming index 0 = High Risk
+        confidence_high_risk = probabilities[0]
         risk_level = "High Risk of Non-Survival" if confidence_high_risk > 0.5 else "Low Risk of Non-Survival"
 
         result = {
@@ -147,4 +155,3 @@ def predict(data: OneHotPatientData):
     except Exception as e:
         logger.error(f"Prediction error: {str(e)}")
         return {"error": str(e)}
-