schemas.py

from pydantic import BaseModel, Field, field_validator
from enum import Enum


class RetinopathyLevel(str, Enum):
    NO_DR = "No DR"
    MILD = "Mild"
    MODERATE = "Moderate"
    SEVERE = "Severe"
    PROLIFERATIVE = "Proliferative DR"
    UNCERTAIN = "Uncertain"


class Retinopathy(BaseModel):
    id: int
    image: bytes 

    @field_validator("image")
    @classmethod
    def check_image_size(cls, v: bytes):
        max_size = 5 * 1024 * 1024  # 5MB

        if len(v) > max_size:
            raise ValueError("Image size must be less than 5MB")

        return v


class DetectionResult(BaseModel):
    id: int
    retinopathy_level: RetinopathyLevel
    confidence: float  

class DiabetesInput(BaseModel):
    Pregnancies: int
    Glucose: int
    BloodPressure: int
    SkinThickness: int
    Insulin: int
    BMI: float
    DiabetesPedigreeFunction: float
    Age: int

class DiabetesOutput(BaseModel):
    prediction: int
    probability: float
    feature_importance: dict
    
class SymptomsInput(BaseModel):
    Age: int = Field(gt=0)
    Gender: int = Field(ge=0, le=1)  # 0=Female, 1=Male
    Polyuria: int = Field(ge=0, le=1)
    Polydipsia: int = Field(ge=0, le=1)
    sudden_weight_loss: int = Field(ge=0, le=1)
    weakness: int = Field(ge=0, le=1)
    Polyphagia: int = Field(ge=0, le=1)
    Genital_thrush: int = Field(ge=0, le=1)
    visual_blurring: int = Field(ge=0, le=1)
    Itching: int = Field(ge=0, le=1)
    Irritability: int = Field(ge=0, le=1)
    delayed_healing: int = Field(ge=0, le=1)
    partial_paresis: int = Field(ge=0, le=1)
    muscle_stiffness: int = Field(ge=0, le=1)
    Alopecia: int = Field(ge=0, le=1)
    Obesity: int = Field(ge=0, le=1)

class SymptomsOutput(BaseModel):
    prediction: int
    probability: float
    feature_importance: dict