kasaam89/Census-Income-Classifier

1from fastapi import FastAPI
2from pydantic import BaseModel
3import pandas as pd
4import pickle
5
6# 1. FastAPI ka app object banayein
7app = FastAPI(title="Adult Income Prediction API")
8
9# 2. Model ko load karein (Jo pickle se freeze kiya tha)
10with open('model.pkl', 'rb') as file:
11    model = pickle.load(file)
12
13# 3. Pydantic Model banayein (Yeh check karega ke frontend se data sahi aa raha hai)
14# Jo columns aap ke dataset mein hain, unhi ke naam yahan honge
15class IncomeInput(BaseModel):
16    age: int
17    workclass: str
18    fnlwgt: int
19    education: str
20    education_num: int  # dhyan rahe columns ke naam dataset jaise hon
21    marital_status: str
22    occupation: str
23    relationship: str
24    race: str
25    gender: str
26    capital_gain: int
27    capital_loss: int
28    hours_per_week: int
29    native_country: str
30
31# 4. Home Route (Sirf check karne ke liye ke API chal rahi hai)
32@app.get("/")
33def home():
34    return {"message": "Adult Income Prediction API is running successfully!"}
35
36# 5. Prediction Route (Yahan frontend data bhejega)
37@app.post("/predict")
38def predict_income(data: IncomeInput):
39    # JSON data ko dictionary mein badlein
40    input_dict = data.model_dump()
41    
42    # ML Pipeline ke columns ke naam exact match karne ke liye '-' lagayein jo dataset mein thay
43    # Kyunke Python variables mein '-' allow nahi karta, isliye hum yahan mapping kar rahe hain
44    formatted_dict = {
45        'age': input_dict['age'],
46        'workclass': input_dict['workclass'],
47        'fnlwgt': input_dict['fnlwgt'],
48        'education': input_dict['education'],
49        'educational-num': input_dict['education_num'],
50        'marital-status': input_dict['marital_status'],
51        'occupation': input_dict['occupation'],
52        'relationship': input_dict['relationship'],
53        'race': input_dict['race'],
54        'gender': input_dict['gender'],
55        'capital-gain': input_dict['capital_gain'],
56        'capital-loss': input_dict['capital_loss'],
57        'hours-per-week': input_dict['hours_per_week'],
58        'native-country': input_dict['native_country']
59    }
60    
61    # Pandas DataFrame banayein kyunke humari pipeline DataFrame leti hai
62    df_input = pd.DataFrame([formatted_dict])
63    
64    # Prediction karein (Pipeline khud encode aur scale karegi)
65    prediction = model.predict(df_input)[0]
66    
67    # Output return karein
68    result = ">50K" if prediction == 1 else "<=50K"
69    return {"prediction": result}