Spaces:
Paused
Paused
import logging.config | |
from enum import Enum | |
from typing import List, Union | |
from pydantic import BaseModel, validator | |
logger = logging.getLogger(__name__) | |
class PredictionResponse(BaseModel): | |
label: str | |
score: float | |
class BedType(str, Enum): | |
none = "" | |
single = "single" | |
double = "double" | |
queen = "queen" | |
king = "king" | |
california_king = "california_king" | |
bunk = "bunk" | |
sofa = "sofa" | |
rollaway = "rollaway" | |
futon = "futon" | |
class BedData(BaseModel): | |
type: Union[BedType, None] = None | |
count: Union[int, None] = None | |
def validate_count(cls, v): | |
return v if v and v > 0 else None | |
def validate_type(cls, v): | |
try: | |
return BedType(v.lower()) if v else None | |
except ValueError: | |
return None | |
def create_valid_bed(cls, type: Union[str, None], count: Union[int, None]): | |
if type and count and count > 0: | |
try: | |
type_enum = BedType(type.lower()) | |
return cls(type=type_enum, count=count) | |
except ValueError: | |
logger.error(f"Invalid bed type: {type}") | |
return None | |
def __init__(self, **data): | |
variations = { | |
"individual": BedType.single, | |
"camaindividual": BedType.single, | |
"doble": BedType.double, | |
"camadoble": BedType.double, | |
"reina": BedType.queen, | |
"queenbed": BedType.queen, | |
"rey": BedType.king, | |
"californiaking": BedType.california_king, | |
"litera": BedType.bunk, | |
"sofá": BedType.sofa, | |
"sofacama": BedType.sofa, | |
"plegable": BedType.rollaway, | |
"futón": BedType.futon, | |
"twin": BedType.single, | |
"twinbed": BedType.single, | |
"singlebed": BedType.single, | |
"doublebed": BedType.double, | |
"largedouble": BedType.queen, | |
"extralargedouble": BedType.king, | |
"bunkbed": BedType.bunk, | |
"couch": BedType.sofa, | |
"airmattress": BedType.futon, | |
"floormattress": BedType.futon, | |
} | |
bed_type = data.get("type") | |
if bed_type: | |
normalized_bed_type = bed_type.replace(" ", "").lower() | |
data["type"] = variations.get(normalized_bed_type, bed_type) | |
super().__init__(**data) | |
class RoomData(BaseModel): | |
room_description: str | |
beds: Union[BedData, None] | |
class Predictions(BaseModel): | |
type: PredictionResponse | |
category: PredictionResponse | |
environment: List[PredictionResponse] | |
feature: List[PredictionResponse] | |
view: List[PredictionResponse] | |
language_detected: str | |
beds: List[BedData] | |
class AllPredictionsResponse(BaseModel): | |
predictions: Predictions | |