Spaces:
Paused
Paused
File size: 2,900 Bytes
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 |
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
@validator("count", pre=True, always=True)
def validate_count(cls, v):
return v if v and v > 0 else None
@validator("type", pre=True, always=True)
def validate_type(cls, v):
try:
return BedType(v.lower()) if v else None
except ValueError:
return None
@classmethod
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
|