Update app.py

app.py

@@ -37,21 +37,45 @@ model_path = hf_hub_download(
 model = joblib.load(model_path)
 print("Model loaded successfully.")
 
-# --- Define Input Data Model ---
-class PatientData(BaseModel):
+# --- Define Input Data Model for one-hot encoded features ---
+# This new model directly matches the one-hot encoded data from the frontend
+class OneHotPatientData(BaseModel):
     age: float
-    gender: str
-    country: str
-    cancer_stage: str
-    family_history: int
-    smoking_status: str
     bmi: float
     cholesterol_level: float
     hypertension: int
     asthma: int
     cirrhosis: int
     other_cancer: int
-    treatment_type: str
+    family_history_Yes: int
+    gender_Male: int
+    gender_Female: int
+    country_Sweden: int
+    country_Netherlands: int
+    country_Hungary: int
+    country_Belgium: int
+    country_Italy: int
+    country_Croatia: int
+    country_Denmark: int
+    country_Germany: int
+    country_France: int
+    country_Slovakia: int
+    country_Finland: int
+    country_Spain: int
+    country_UnitedKingdom: int
+    country_UnitedStates: int
+    cancer_stage_StageI: int
+    cancer_stage_StageII: int
+    cancer_stage_StageIII: int
+    cancer_stage_StageIV: int
+    smoking_status_NeverSmoked: int
+    smoking_status_FormerSmoker: int
+    smoking_status_PassiveSmoker: int
+    smoking_status_CurrentSmoker: int
+    treatment_type_Chemotherapy: int
+    treatment_type_Surgery: int
+    treatment_type_Radiation: int
+    treatment_type_Combined: int
 
 # --- API Endpoints ---
 @app.get("/")
@@ -59,67 +83,9 @@ def read_root():
     return {"message": "Welcome to the PulmoProbe AI API"}
 
 @app.post("/predict")
-def predict(data: PatientData):
+def predict(data: OneHotPatientData):
     try:
-        # Define the exact list of features your model expects
-        feature_order = [
-            'age', 'bmi', 'cholesterol_level', 'hypertension', 'asthma', 'cirrhosis',
-            'other_cancer', 'gender_Female', 'gender_Male',
-            'country_Belgium', 'country_Croatia', 'country_Denmark', 'country_Finland', 
-            'country_France', 'country_Germany', 'country_Hungary', 'country_Italy',
-            'country_Netherlands', 'country_Slovakia', 'country_Spain', 'country_Sweden',
-            'cancer_stage_Stage I', 'cancer_stage_Stage II', 'cancer_stage_Stage III', 
-            'cancer_stage_Stage IV',
-            'family_history_Yes',
-            'smoking_status_Current Smoker', 'smoking_status_Former Smoker',
-            'smoking_status_Never Smoked', 'smoking_status_Passive Smoker',
-            'treatment_type_Chemotherapy', 'treatment_type_Combined', 
-            'treatment_type_Radiation', 'treatment_type_Surgery'
-        ]
-
-        # Convert the Pydantic model to a dictionary
-        input_data = data.dict()
-        
-        # Initialize an empty dictionary for one-hot encoded features
-        encoded_data = {feature: 0 for feature in feature_order}
-        
-        # Map original data to the one-hot encoded format
-        encoded_data['age'] = input_data['age']
-        encoded_data['bmi'] = input_data['bmi']
-        encoded_data['cholesterol_level'] = input_data['cholesterol_level']
-        encoded_data['hypertension'] = input_data['hypertension']
-        encoded_data['asthma'] = input_data['asthma']
-        encoded_data['cirrhosis'] = input_data['cirrhosis']
-        encoded_data['other_cancer'] = input_data['other_cancer']
-
-        # One-hot encode categorical features
-        gender_key = f"gender_{input_data['gender']}"
-        if gender_key in encoded_data:
-            encoded_data[gender_key] = 1
-
-        country_key = f"country_{input_data['country']}"
-        if country_key in encoded_data:
-            encoded_data[country_key] = 1
-            
-        cancer_stage_key = f"cancer_stage_{input_data['cancer_stage']}"
-        if cancer_stage_key in encoded_data:
-            encoded_data[cancer_stage_key] = 1
-
-        smoking_status_key = f"smoking_status_{input_data['smoking_status']}"
-        if smoking_status_key in encoded_data:
-            encoded_data[smoking_status_key] = 1
-
-        treatment_type_key = f"treatment_type_{input_data['treatment_type']}"
-        if treatment_type_key in encoded_data:
-            encoded_data[treatment_type_key] = 1
-            
-        # Family history is already 0/1, just rename the key
-        encoded_data['family_history_Yes'] = input_data['family_history']
-        
-        # Create the DataFrame with the correct order
-        input_df = pd.DataFrame([encoded_data])
-        
-        # Make the prediction
+        input_df = pd.DataFrame([data.dict()])
         probabilities = model.predict_proba(input_df)[0]
         confidence_high_risk = probabilities[0]
         risk_level = "High Risk of Non-Survival" if confidence_high_risk > 0.5 else "Low Risk of Non-Survival"
@@ -129,5 +95,4 @@ def predict(data: PatientData):
             "confidence": f"{confidence_high_risk * 100:.1f}%"
         }
     except Exception as e:
-        # Return a detailed error message for better debugging
-        return {"error": str(e), "input_data_received": data.dict()}
+        return {"error": str(e)}