Update app.py

app.py

@@ -9,22 +9,66 @@ from huggingface_hub import hf_hub_download
 import os
 import logging
 
-# ... (logging setup, FastAPI initialization, and model loading remain the same) ...
+# ------------------------------------------------------------
+# Setup Logging
+# ------------------------------------------------------------
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # ------------------------------------------------------------
-# Define Input Schema (Corrected Names)
+# FastAPI Initialization
+# ------------------------------------------------------------
+# This line MUST come before any @app.<method> decorators
+app = FastAPI(title="PulmoProbe AI API")
+
+# Allow CORS for frontend communication
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Use specific domain in production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ------------------------------------------------------------
+# Hugging Face Model Setup
+# ------------------------------------------------------------
+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']}")
+
+MODEL_REPO_ID = "costaspinto/PulmoProbe"
+MODEL_FILENAME = "best_model.joblib"
+
+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 load model: {str(e)}")
+    raise RuntimeError(f"Model loading failed: {str(e)}")
+
+# ------------------------------------------------------------
+# Define Input Schema
 # ------------------------------------------------------------
 class OneHotPatientData(BaseModel):
+    # Continuous fields
     age: float
     bmi: float
     cholesterol_level: float
+    # Binary medical conditions
     hypertension: int
     asthma: int
     cirrhosis: int
     other_cancer: int
+    # Gender (Male = 1, Female = 0)
     gender_Male: int
-    family_history_Yes: int
-    
+    # Countries (One-Hot)
     country_Belgium: int
     country_Bulgaria: int
     country_Croatia: int
@@ -51,23 +95,30 @@ class OneHotPatientData(BaseModel):
     country_Slovenia: int
     country_Spain: int
     country_Sweden: int
-
-    cancer_stage_Stage_II: int
-    cancer_stage_Stage_III: int
-    cancer_stage_Stage_IV: 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 remains the same) ...
+# ------------------------------------------------------------
+# Root Endpoint
+# ------------------------------------------------------------
+@app.get("/")
+def read_root():
+    return {"message": "Welcome to the PulmoProbe AI API"}
 
 # ------------------------------------------------------------
-# Prediction Endpoint (Corrected Feature Order)
+# Prediction Endpoint
 # ------------------------------------------------------------
 @app.post("/predict")
 def predict(data: OneHotPatientData):
@@ -75,26 +126,26 @@ def predict(data: OneHotPatientData):
         input_dict = data.dict()
         logger.info(f"Incoming data: {input_dict}")
         
-        # Define the exact feature order your model expects (with underscores)
+        # Define the exact feature order your model expects
         feature_order = [
             'age', 'bmi', 'cholesterol_level', 'hypertension', 'asthma', 
             'cirrhosis', 'other_cancer', 'gender_Male', 'country_Belgium', 
             'country_Bulgaria', 'country_Croatia', 'country_Cyprus', 
-            'country_Czech_Republic', 'country_Denmark', 'country_Estonia', 
+            'country_Czech Republic', 'country_Denmark', 'country_Estonia', 
             'country_Finland', 'country_France', 'country_Germany', 
             'country_Greece', 'country_Hungary', 'country_Ireland', 
             'country_Italy', 'country_Latvia', 'country_Lithuania', 
             'country_Luxembourg', 'country_Malta', 'country_Netherlands', 
             'country_Poland', 'country_Portugal', 'country_Romania', 
             'country_Slovakia', 'country_Slovenia', 'country_Spain', 
-            'country_Sweden', 'cancer_stage_Stage_II', 'cancer_stage_Stage_III', 
-            'cancer_stage_Stage_IV', 'family_history_Yes', 
-            'smoking_status_Former_Smoker', 'smoking_status_Never_Smoked', 
-            'smoking_status_Passive_Smoker', 'treatment_type_Combined', 
+            'country_Sweden', 'cancer_stage_Stage Ii', 'cancer_stage_Stage Iii', 
+            'cancer_stage_Stage Iv', 'family_history_Yes', 
+            'smoking_status_Former Smoker', 'smoking_status_Never Smoked', 
+            'smoking_status_Passive Smoker', 'treatment_type_Combined', 
             'treatment_type_Radiation', 'treatment_type_Surgery'
         ]
 
-        # Create DataFrame and ensure the columns are in the correct order
+        # Convert dictionary to a DataFrame and ensure the columns are in the correct order
         input_df = pd.DataFrame([input_dict], columns=feature_order)
         logger.info(f"DataFrame for prediction: {input_df}")
 
@@ -102,6 +153,7 @@ def predict(data: OneHotPatientData):
         probabilities = model.predict_proba(input_df)[0]
         logger.info(f"Model probabilities: {probabilities}")
 
+        # 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"
 
@@ -115,4 +167,4 @@ def predict(data: OneHotPatientData):
 
     except Exception as e:
         logger.error(f"Prediction error: {str(e)}")
-        return {"error": str(e)}
+        return {"error": str(e), "input_data_received": data.dict()}