Update prediction.py

Author: codewithdark-git

app/schemas/prediction.py

@@ -1,21 +1,17 @@
-from pydantic import BaseModel, Field
-from typing import Literal, Optional
 from enum import Enum
+from pydantic import BaseModel
 
 class DiseaseType(str, Enum):
-    PNEUMONIA = "Pneumonia"
-    MELANOMA = "Melanoma"
-    BRAIN = "Brain"
-    HEART = "Heart"
-
-class PredictionRequest(BaseModel):
-    disease_type: DiseaseType = Field(..., description="Type of disease to predict")
-    image_base64: str = Field(..., description="Base64 encoded image data")
+    SKIN_CANCER = "skin_cancer"
+    PNEUMONIA = "pneumonia"
+    BRAIN_TUMOR = "brain_tumor"
+    EYE_DISEASE = "eye_disease"
+    HEART_DISEASE = "heart_disease"
 
 class PredictionResponse(BaseModel):
     disease_type: DiseaseType
     prediction: str
-    confidence: float = Field(..., ge=0.0, le=1.0)
+    confidence: float
     analysis: str
-    status: Literal["success", "error"] = "success"
-    error_message: Optional[str] = None 
+    status: str = "success"
+    error_message: str | None = None