update main

main.py

@@ -1,8 +1,7 @@
 import uvicorn
 import pandas as pd
-from pydantic import BaseModel
-from typing import List, Union
-from fastapi import FastAPI
+from typing import Union
+from fastapi import FastAPI, Query
 import joblib
 from enum import Enum
 from fastapi.responses import HTMLResponse
@@ -56,20 +55,33 @@ class CarBrand(str, Enum):
     subaru = "Subaru"
     toyota = "Toyota"
 
-class PredictionFeatures(BaseModel):
-    brand: CarBrand
-    mileage: int
-    engine_power: int
-    fuel: str
-    paint_color: str
-    car_type: str
-    private_parking_available: bool
-    has_gps: bool
-    has_air_conditioning: bool
-    automatic_car: bool
-    has_getaround_connect: bool
-    has_speed_regulator: bool
-    winter_tires: bool
+class FuelType(str, Enum):
+    diesel = "diesel"
+    petrol = "petrol"
+    hybrid_petrol = "hybrid_petrol"
+    electro = "electro"
+
+class PaintColor(str, Enum):
+    black = "black"
+    grey = "grey"
+    white = "white"
+    red = "red"
+    silver = "silver"
+    blue = "blue"
+    orange = "orange"
+    beige = "beige"
+    brown = "brown"
+    green = "green"
+
+class CarType(str, Enum):
+    convertible = "convertible"
+    coupe = "coupe"
+    estate = "estate"
+    hatchback = "hatchback"
+    sedan = "sedan"
+    subcompact = "subcompact"
+    suv = "suv"
+    van = "van"
 
 @app.get("/", response_class=HTMLResponse, tags=["Introduction Endpoints"])
 async def index():
@@ -81,15 +93,37 @@ async def index():
     )
 
 @app.post("/predict", tags=["Machine Learning"])
-async def predict(predictionFeatures: PredictionFeatures):
-    columns = [
-        'brand', 'mileage', 'engine_power', 'fuel', 'paint_color',
-        'car_type', 'private_parking_available', 'has_gps',
-        'has_air_conditioning', 'automatic_car', 'has_getaround_connect',
-        'has_speed_regulator', 'winter_tires'
-    ]
+async def predict(
+    brand: CarBrand,
+    mileage: int = Query(...),
+    engine_power: int = Query(...),
+    fuel: FuelType = Query(...),
+    paint_color: PaintColor = Query(...),
+    car_type: CarType = Query(...),
+    private_parking_available: bool = Query(...),
+    has_gps: bool = Query(...),
+    has_air_conditioning: bool = Query(...),
+    automatic_car: bool = Query(...),
+    has_getaround_connect: bool = Query(...),
+    has_speed_regulator: bool = Query(...),
+    winter_tires: bool = Query(...)
+):
     
-    car_data_dict = {col: [getattr(predictionFeatures, col)] for col in columns}
+    car_data_dict = {
+        'model_key': [brand],
+        'mileage': [mileage],
+        'engine_power': [engine_power],
+        'fuel': [fuel],
+        'paint_color': [paint_color],
+        'car_type': [car_type],
+        'private_parking_available': [private_parking_available],
+        'has_gps': [has_gps],
+        'has_air_conditioning': [has_air_conditioning],
+        'automatic_car': [automatic_car],
+        'has_getaround_connect': [has_getaround_connect],
+        'has_speed_regulator': [has_speed_regulator],
+        'winter_tires': [winter_tires]
+    }
     car_data = pd.DataFrame(car_data_dict)
 
     model = joblib.load('best_model_XGBoost.pkl')