Prepare job build context

app.py

@@ -66,13 +66,22 @@ def generate_slimsam_mask(rgb_image: np.ndarray, box_prompt: List[int]) -> tuple
 
     model, processor, device = _get_slimsam()
     raw_image = Image.fromarray(rgb_image).convert("RGB")
-    inputs = processor(raw_image, input_boxes=[[box_prompt]], return_tensors="pt").to(device)
+    enc = processor(raw_image, input_boxes=[[box_prompt]], return_tensors="np")
+    # Keep size tensors on CPU for post-processing
+    original_sizes = torch.as_tensor(enc["original_sizes"])
+    reshaped_sizes = torch.as_tensor(enc["reshaped_input_sizes"])
+    # Move model inputs to device
+    inputs = {
+        k: torch.as_tensor(v).to(device)
+        for k, v in enc.items()
+        if k not in {"original_sizes", "reshaped_input_sizes"}
+    }
     outputs = model(**inputs)
 
     masks = processor.image_processor.post_process_masks(
         outputs.pred_masks.cpu(),
-        inputs["original_sizes"].cpu(),
-        inputs["reshaped_input_sizes"].cpu(),
+        original_sizes,
+        reshaped_sizes,
     )[0]
     scores = outputs.iou_scores.squeeze().cpu()
     best_idx = int(scores.argmax().item())
@@ -533,11 +542,11 @@ with gr.Blocks(title="FoundationPose Inference", theme=gr.themes.Soft()) as demo
 
                             gr.Markdown("### Camera Intrinsics")
                             with gr.Row():
-                                cad_fx = gr.Number(label="fx", value=500.0)
-                                cad_fy = gr.Number(label="fy", value=500.0)
+                                cad_fx = gr.Number(label="fx", value=193.13708498984758)
+                                cad_fy = gr.Number(label="fy", value=193.13708498984758)
                             with gr.Row():
-                                cad_cx = gr.Number(label="cx", value=320.0)
-                                cad_cy = gr.Number(label="cy", value=240.0)
+                                cad_cx = gr.Number(label="cx", value=120.0)
+                                cad_cy = gr.Number(label="cy", value=80.0)
 
                             cad_init_button = gr.Button("Initialize with CAD", variant="primary")
 
@@ -591,11 +600,11 @@ with gr.Blocks(title="FoundationPose Inference", theme=gr.themes.Soft()) as demo
 
                     gr.Markdown("### Camera Intrinsics")
                     with gr.Row():
-                        est_fx = gr.Number(label="fx (focal length x)", value=500.0)
-                        est_fy = gr.Number(label="fy (focal length y)", value=500.0)
+                        est_fx = gr.Number(label="fx (focal length x)", value=193.13708498984758)
+                        est_fy = gr.Number(label="fy (focal length y)", value=193.13708498984758)
                     with gr.Row():
-                        est_cx = gr.Number(label="cx (principal point x)", value=320.0)
-                        est_cy = gr.Number(label="cy (principal point y)", value=240.0)
+                        est_cx = gr.Number(label="cx (principal point x)", value=120.0)
+                        est_cy = gr.Number(label="cy (principal point y)", value=80.0)
 
                     est_button = gr.Button("Estimate Pose", variant="primary")
 

estimator.py

@@ -300,6 +300,9 @@ class FoundationPoseEstimator:
         Returns:
             Dictionary with position, orientation (quaternion), and confidence
         """
+        if torch.is_tensor(pose_matrix):
+            pose_matrix = pose_matrix.detach().cpu().numpy()
+
         # Extract translation
         translation = pose_matrix[:3, 3]