Update app/models.py

app/models.py

@@ -1,7 +1,11 @@
+# app/models.py
 from pydantic import BaseModel
-from typing import Optional
+from typing import Optional, List
 
 
+# =========================
+# Upload
+# =========================
 class UploadResponse(BaseModel):
     image_id: str
     filename: str
@@ -11,19 +15,24 @@ class UploadResponse(BaseModel):
     url: str
 
 
+# =========================
+# Classification
+# =========================
 class ClassifyRequest(BaseModel):
     image_id: str
-    model_version: Optional[str] = None
 
 
 class ClassifyResponse(BaseModel):
     image_id: str
     prediction: str
     confidence: float
-    model_version: str
+    model_version: str = "hf_savedmodel"
     status: str = "completed"
 
 
+# =========================
+# Segmentation
+# =========================
 class SegmentRequest(BaseModel):
     image_id: str
 
@@ -31,7 +40,7 @@ class SegmentRequest(BaseModel):
 class SegmentResponse(BaseModel):
     image_id: str
     status: str
-    masks_shape: Optional[list] = None
+    masks_shape: Optional[List[int]] = None
     max_confidence: Optional[float] = None
     result_url: Optional[str] = None
-    error: Optional[str] = None
+    error: Optional[str] = None