K as param

app.py

@@ -52,7 +52,7 @@ def depth_to_normal(depth):
 # ===============================
 # CORE PROCESSING FUNCTION
 # ===============================
-def _process_saree_core(base_image: Image.Image, pattern_image: Image.Image):
+def _process_saree_core(base_image: Image.Image, pattern_image: Image.Image, k: int = 5):
     img_pil = base_image.convert("RGB")
     img_np = np.array(img_pil)
 
@@ -123,8 +123,7 @@ def _process_saree_core(base_image: Image.Image, pattern_image: Image.Image):
     # Extract alpha and clean edges
     mask_alpha = np.array(mask_img)[:, :, 3].astype(np.float32) / 255.0
 
-    # 1. Slightly stronger shrink (balanced)
-    k = 8
+    # 1. Slightly stronger shrink (balanced)    
     kernel = np.ones((k, k), np.uint8)   # slightly larger kernel
     mask_binary = (mask_alpha > k/100).astype(np.uint8) * 255  # slightly stricter threshold
     mask_eroded = cv2.erode(mask_binary, kernel, iterations=3)  # balanced erosion
@@ -149,7 +148,7 @@ def _process_saree_core(base_image: Image.Image, pattern_image: Image.Image):
 # ===============================
 # WRAPPER: ACCEPT BYTES OR BASE64
 # ===============================
-def process_saree(data):
+def process_saree(data, k: int = 5):
     if not isinstance(data, (list, tuple)) or len(data) != 2:
         raise HTTPException(status_code=422, detail="Expected an array with two elements: [base_blob, pattern_blob]")
 
@@ -170,7 +169,7 @@ def process_saree(data):
     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)
+    return _process_saree_core(base_image, pattern_image, k=k)
 
 # ===============================
 # GRADIO + FASTAPI APP
@@ -198,31 +197,29 @@ async def root():
 # Mount Gradio at /gradio
 app = gr.mount_gradio_app(app, gradio_iface, path="/gradio")
 
+# ===============================
 # Custom API endpoint
+# ===============================
 @app.post("/predict-saree")
 async def predict_saree(request: Request):
     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"])
+        k = body.get("k", 5)  # Default to 5 if not provided
+        result_img = process_saree(body["data"], k=int(k))
 
         buf = BytesIO()
         result_img.save(buf, format="PNG")
         base64_img = base64.b64encode(buf.getvalue()).decode("utf-8")
-
         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}
-        )
+        return JSONResponse(status_code=500, content={"error": "Processing Error", "details": str(e), "trace": tb})
 
 # Alias for backward compatibility
 @app.post("/api/predict/")