Update app.py

app.py

@@ -14,7 +14,7 @@ def load_description(fp):
 
 
 @spaces.GPU(duration=90)
-def run_theia(image):
+def run_theia(image, pred_iou_thresh, stability_score_thresh):
     theia_model = AutoModel.from_pretrained("theaiinstitute/theia-tiny-patch16-224-cddsv", trust_remote_code=True)
     theia_model = theia_model.to('cuda')
     target_model_names = [
@@ -41,19 +41,22 @@ def run_theia(image):
         mask_generator=mask_generator,
         sam_model=sam_model,
         depth_anything_decoder=depth_anything_decoder,
-        pred_iou_thresh=0.5,
-        stability_score_thresh=0.7,
+        pred_iou_thresh=pred_iou_thresh,
+        stability_score_thresh=stability_score_thresh,
         gt=True,
         device='cuda',
     )
 
     _, width, _ = theia_decode_results[0].shape
-    theia_decode_dino = (255.0 * theia_decode_results[0]).astype(np.uint8)[:, width // 4 : 2 * width // 4, :]
-    theia_decode_sam = (255.0 * theia_decode_results[0]).astype(np.uint8)[:, 2 * width // 4 : 3 * width // 4, :]
-    theia_decode_depth = (255.0 * theia_decode_results[0]).astype(np.uint8)[:, 3 * width // 4 :, :]
-    gt_dino = (255.0 * gt_results[0]).astype(np.uint8)[:, width // 4 : 2 * width // 4, :]
-    gt_sam = (255.0 * gt_results[0]).astype(np.uint8)[:, 2 * width // 4 : 3 * width // 4, :]
-    gt_depth = (255.0 * gt_results[0]).astype(np.uint8)[:, 3 * width // 4 :, :]
+    theia_decode_results = (255.0 * theia_decode_results[0]).astype(np.uint8)
+    theia_decode_dino = theia_decode_results[:, width // 4 : 2 * width // 4, :]
+    theia_decode_sam = theia_decode_results[:, 2 * width // 4 : 3 * width // 4, :]
+    theia_decode_depth = theia_decode_results[:, 3 * width // 4 :, :]
+
+    gt_results = (255.0 * gt_results[0]).astype(np.uint8)
+    gt_dino = gt_results[:, width // 4 : 2 * width // 4, :]
+    gt_sam = gt_results[:, 2 * width // 4 : 3 * width // 4, :]
+    gt_depth = gt_results[:, 3 * width // 4 :, :]
 
     dinov2_output = [(theia_decode_dino, "Theia"), (gt_dino, "Ground Truth")]
     sam_output = [(theia_decode_sam, "Theia"), (gt_sam, "Ground Truth")]
@@ -67,6 +70,8 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(label="Input Image", type="pil")
+            pred_iou_thresh = gr.Slider(0.0, 1.0, value=0.5, label="Pred IoU Thresh")
+            stability_score_thresh = gr.Slider(0.0, 1.0, value=0.7, label="Stability Score Thresh")
             submit_button = gr.Button("Submit")
         
         with gr.Column():
@@ -74,7 +79,11 @@ with gr.Blocks() as demo:
             sam_output = gr.Gallery(label="SAM", type="numpy")
             depth_anything_output = gr.Gallery(label="Depth-Anything", type="numpy")
 
-    submit_button.click(run_theia, inputs=input_image, outputs=[dinov2_output, sam_output, depth_anything_output])
+    submit_button.click(
+        run_theia,
+        inputs=[input_image, pred_iou_thresh, stability_score_thresh],
+        outputs=[dinov2_output, sam_output, depth_anything_output]
+    )
 
 demo.queue()
 demo.launch()