Prepare job build context

app.py

@@ -352,7 +352,8 @@ def gradio_estimate(
     fy: float,
     cx: float,
     cy: float,
-    mask_method: str
+    mask_method: str,
+    mask_editor_data
 ):
     """Gradio wrapper for pose estimation."""
     try:
@@ -409,6 +410,33 @@ def gradio_estimate(
             logger.info("Otsu mask coverage %.1f%%", mask_percentage)
             if fallback_full_image:
                 logger.warning("Otsu mask fallback to full image due to unrealistic coverage")
+        elif mask_method == "From editor":
+            editor_mask = None
+            if isinstance(mask_editor_data, dict):
+                layers = mask_editor_data.get("layers")
+                if isinstance(layers, list) and layers:
+                    editor_mask = layers[-1]
+                else:
+                    editor_mask = mask_editor_data.get("composite")
+            else:
+                editor_mask = mask_editor_data
+
+            if editor_mask is None:
+                return "Error: No editor mask provided", query_image, None
+
+            editor_mask = np.array(editor_mask)
+            if editor_mask.ndim == 3 and editor_mask.shape[2] >= 4:
+                alpha = editor_mask[:, :, 3]
+                mask = (alpha > 0).astype(np.uint8) * 255
+            elif editor_mask.ndim == 3:
+                gray = cv2.cvtColor(editor_mask, cv2.COLOR_RGB2GRAY)
+                mask = (gray > 0).astype(np.uint8) * 255
+            elif editor_mask.ndim == 2:
+                mask = (editor_mask > 0).astype(np.uint8) * 255
+            else:
+                return "Error: Unsupported editor mask format", query_image, None
+
+            debug_mask = mask
 
         # Estimate pose
         result = pose_estimator.estimate_pose(
@@ -593,11 +621,17 @@ with gr.Blocks(title="FoundationPose Inference", theme=gr.themes.Soft()) as demo
                     )
 
                     est_mask_method = gr.Radio(
-                        choices=["SlimSAM", "Otsu"],
+                        choices=["SlimSAM", "Otsu", "From editor"],
                         value="SlimSAM",
                         label="Mask Method"
                     )
 
+                    est_mask_editor = gr.ImageEditor(
+                        label="Mask Editor (paint mask)",
+                        type="numpy",
+                        visible=False
+                    )
+
                     est_fx = gr.Number(label="fx (focal length x)", value=193.13708498984758, visible=False)
                     est_fy = gr.Number(label="fy (focal length y)", value=193.13708498984758, visible=False)
                     est_cx = gr.Number(label="cx (principal point x)", value=120.0, visible=False)
@@ -614,9 +648,28 @@ with gr.Blocks(title="FoundationPose Inference", theme=gr.themes.Soft()) as demo
                     )
                     est_viz = gr.Image(label="Query Image")
 
+            def _toggle_editor(method: str):
+                return gr.update(visible=method == "From editor")
+
+            est_mask_method.change(
+                fn=_toggle_editor,
+                inputs=est_mask_method,
+                outputs=est_mask_editor
+            )
+
             est_button.click(
                 fn=gradio_estimate,
-                inputs=[est_object_id, est_query_image, est_depth_image, est_fx, est_fy, est_cx, est_cy, est_mask_method],
+                inputs=[
+                    est_object_id,
+                    est_query_image,
+                    est_depth_image,
+                    est_fx,
+                    est_fy,
+                    est_cx,
+                    est_cy,
+                    est_mask_method,
+                    est_mask_editor,
+                ],
                 outputs=[est_output, est_viz, est_mask]
             )