Spaces:
Paused
Paused
File size: 6,809 Bytes
f006f31 7b68836 f006f31 7b68836 f006f31 7b68836 f006f31 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 |
import logging
import time
from concurrent.futures import ThreadPoolExecutor
from fastapi import APIRouter, Query, Request
from mappingservice.constants import DEFAULT_LABEL, DEFAULT_SCORE
from mappingservice.dependencies import (
mc,
)
from mappingservice.models import (
AllPredictionsResponse,
PredictionResponse,
Predictions,
RoomData,
)
from mappingservice.ms.ml_models.bed_type import BedType as BedTypeModel
from mappingservice.ms.ml_models.environment import Environment
from mappingservice.ms.ml_models.room_category import RoomCategory
from mappingservice.ms.ml_models.room_features import RoomFeatures
from mappingservice.ms.ml_models.room_type import RoomType
from mappingservice.ms.ml_models.room_view import RoomView
from mappingservice.utils import (
get_bed_predictions,
process_predictions,
safe_round,
)
logging.basicConfig(
level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)
router = APIRouter(
prefix="/predict/room", tags=["room"], responses={404: {"description": "Not found"}}
)
def get_room_type_prediction(room_description: str, language: str = "en"):
pipeline = mc['room_type'][language]
model = BedTypeModel()
return {"msg": model.predict(room_description, pipeline, language)}
def get_view_prediction(room_description: str, language: str = "en"):
pipeline = mc['room_view'][language]
model = RoomView()
return {"view_prediction": model.predict(room_description, pipeline, language)}
def get_room_category_prediction(room_description: str, language: str = "en"):
pipeline = mc['room_category'][language]
model = RoomCategory()
return {"msg": model.predict(room_description, pipeline, language)}
def get_feature_prediction(room_description: str, language: str = "en"):
pipeline = mc['room_features'][language]
model = RoomFeatures()
return {"feature_prediction": model.predict(room_description, pipeline, language)}
def get_room_environment_prediction(room_description: str, language: str = "en"):
pipeline = mc['environment'][language]
model = Environment()
return {"msg": model.predict(room_description, pipeline, language)}
@router.post("/predict/beds")
async def predict_beds(request: Request, room_description: str = Query(...)): # noqa: E501
language = request.state.predicted_language
pipeline = mc['bed_type'][language]
model = BedTypeModel()
prediction = model.predict(room_description, pipeline, language)
return prediction
@router.get("/type")
async def predict_room_type_endpoint(request: Request, room_description: str = Query(...)): # noqa: E501
language = request.state.predicted_language
pipeline = mc['room_type'][language]
model = RoomType()
prediction = model.predict(room_description, pipeline, language)
return prediction
@router.get("/category")
async def predict_room_category_endpoint(request: Request, room_description: str = Query(...)): # noqa: E501
prediction = mc['room_category']['en'].predict(room_description)
return prediction
@router.get("/environment")
async def predict_room_environment_endpoint(request: Request, room_description: str = Query(...)): # noqa: E501
prediction = mc['environment']['en'].predict(room_description)
return prediction
@router.get("/view")
async def predict_view_endpoint(request: Request, room_description: str = Query(...)): # noqa: E501
prediction = mc['room_view']['en'].predict(room_description)
return prediction
@router.get("/feature")
async def predict_feature_endpoint(request: Request, room_description: str = Query(...)): # noqa: E501
prediction = mc['room_features']['en'].predict(room_description)
return prediction
@router.post("/predict/all", response_model=AllPredictionsResponse)
async def predict_all(request: Request, room_data: RoomData):
start_time = time.time()
room_data = RoomData(**await request.json())
language = request.state.predicted_language
with ThreadPoolExecutor() as executor:
type_future = executor.submit(
get_room_type_prediction, room_data.room_description, language
)
category_future = executor.submit(
get_room_category_prediction, room_data.room_description, language
)
environment_future = executor.submit(
get_room_environment_prediction, room_data.room_description, language
)
feature_future = executor.submit(
get_feature_prediction, room_data.room_description, language
)
view_future = executor.submit(
get_view_prediction, room_data.room_description, language
)
type_pred = type_future.result()["msg"]
category_pred = category_future.result()["msg"]
environment_pred_results = environment_future.result()["msg"]
feature_pred_results = feature_future.result()["feature_prediction"]
view_pred_results = view_future.result()["view_prediction"]
bed_predictions = room_data.beds
if not room_data.beds:
logger.debug("No bed data provided or valid; extracting from description.")
extracted_beds = get_bed_predictions(room_data.room_description)
if extracted_beds:
bed_predictions.extend(extracted_beds)
if not isinstance(bed_predictions, list):
bed_predictions = [bed_predictions]
end_time = time.time()
total_time = end_time - start_time
logger.info(f"Total processing time: {total_time:.3f} seconds")
formatted_predictions = {
"type": {
"label": type_pred.get("label", DEFAULT_LABEL),
"score": safe_round(type_pred.get("score", DEFAULT_SCORE), 3),
},
"category": {
"label": category_pred.get("label", DEFAULT_LABEL),
"score": safe_round(category_pred.get("score", DEFAULT_SCORE), 3),
},
}
env_preds = process_predictions(environment_pred_results)
feat_preds = process_predictions(
feature_pred_results.get("features", []), label_key="word"
)
view_preds = process_predictions(
view_pred_results.get("views", []), label_key="word"
)
predictions = Predictions(
type=PredictionResponse(**formatted_predictions["type"]),
category=PredictionResponse(**formatted_predictions["category"]),
environment=[PredictionResponse(**pred) for pred in env_preds] if env_preds else [], # noqa: E501
feature=[PredictionResponse(**pred) for pred in feat_preds] if feat_preds else [], # noqa: E501
view=[PredictionResponse(**pred) for pred in view_preds] if view_preds else [],
language_detected=language,
beds=bed_predictions,
)
return AllPredictionsResponse(predictions=predictions)
|