Error Handling Updated

app.py

@@ -1,15 +1,15 @@
 import gradio as gr
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, HTTPException
 from fastapi.responses import JSONResponse
 import torch
 import torch.nn as nn
 import timm
 import cv2
 import numpy as np
-from PIL import Image
+from PIL import Image, UnidentifiedImageError
 from io import BytesIO
 import base64
-import os
+import traceback
 
 # ===============================
 # SIMPLE DPT MODEL (DEPTH ESTIMATION)
@@ -51,7 +51,6 @@ def depth_to_normal(depth):
 # CORE PROCESSING FUNCTION
 # ===============================
 def _process_saree_core(base_image: Image.Image, pattern_image: Image.Image):
-    # (Depth estimation + pattern blending logic — unchanged)
     img_pil = base_image.convert("RGB")
     img_np = np.array(img_pil)
 
@@ -146,18 +145,25 @@ def process_saree(data):
     Accepts [base_blob, pattern_blob] as bytes OR base64 strings
     """
     if not isinstance(data, (list, tuple)) or len(data) != 2:
-        raise ValueError("Expected an array with two elements: [base_blob, pattern_blob]")
+        raise HTTPException(status_code=422, detail="Expected an array with two elements: [base_blob, pattern_blob]")
 
-    base_blob, pattern_blob = data
+    try:
+        base_blob, pattern_blob = data
 
-    # Convert base64 to bytes if needed
-    if isinstance(base_blob, str):
-        base_blob = base64.b64decode(base_blob.split(",")[-1])
-    if isinstance(pattern_blob, str):
-        pattern_blob = base64.b64decode(pattern_blob.split(",")[-1])
+        # Convert base64 to bytes if needed
+        if isinstance(base_blob, str):
+            base_blob = base64.b64decode(base_blob.split(",")[-1])
+        if isinstance(pattern_blob, str):
+            pattern_blob = base64.b64decode(pattern_blob.split(",")[-1])
 
-    base_image = Image.open(BytesIO(base_blob)).convert("RGBA")
-    pattern_image = Image.open(BytesIO(pattern_blob)).convert("RGBA")
+        base_image = Image.open(BytesIO(base_blob)).convert("RGBA")
+        pattern_image = Image.open(BytesIO(pattern_blob)).convert("RGBA")
+
+    except (base64.binascii.Error, UnidentifiedImageError) as e:
+        raise HTTPException(status_code=422, detail=f"Invalid image data: {str(e)}")
+
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Error reading input images: {str(e)}")
 
     return _process_saree_core(base_image, pattern_image)
 
@@ -180,15 +186,30 @@ app = gr.mount_gradio_app(app, gradio_iface, path="/")
 # Custom named API endpoint
 @app.post("/predict-saree")
 async def predict_saree(request: Request):
-    body = await request.json()
-    result_img = process_saree(body["data"])
+    try:
+        body = await request.json()
+
+        if "data" not in body:
+            raise HTTPException(status_code=422, detail="Missing 'data' field in request body")
+
+        result_img = process_saree(body["data"])
 
-    # Convert output image to base64 PNG
-    buf = BytesIO()
-    result_img.save(buf, format="PNG")
-    base64_img = base64.b64encode(buf.getvalue()).decode("utf-8")
+        # Convert output image to base64 PNG
+        buf = BytesIO()
+        result_img.save(buf, format="PNG")
+        base64_img = base64.b64encode(buf.getvalue()).decode("utf-8")
 
-    return JSONResponse(content={"image_base64": base64_img})
+        return JSONResponse(content={"image_base64": base64_img})
+
+    except HTTPException as e:
+        return JSONResponse(status_code=e.status_code, content={"error": "Input Error", "details": e.detail})
+
+    except Exception as e:
+        tb = traceback.format_exc()
+        return JSONResponse(
+            status_code=500,
+            content={"error": "Processing Error", "details": str(e), "trace": tb}
+        )
 
 # Run (Hugging Face will call uvicorn automatically)
 if __name__ == "__main__":