updated app

app.py

@@ -12,17 +12,14 @@ processor = SamProcessor.from_pretrained('facebook/sam-vit-base')
 model = SamModel.from_pretrained('kitooo/sidewalk-seg-base')
 model.to(device)
 
-def segment_sidewalk(image, threshold):
+def segment_sidewalk(image):
     width, height = image.size
     prompt = [0, 0, width, height]
     inputs = processor(image, input_boxes=[[prompt]], return_tensors='pt')
-    # outputs = model(pixel_values=inputs['pixel_values'].to(device),
-    #                 input_boxes=inputs['input_boxes'].to(device),
-    #                 multimask_output=False)
     with torch.no_grad():
         outputs = model(**inputs, multimask_output=False)
     prob_map = torch.sigmoid(outputs.pred_masks.squeeze()).cpu().detach()
-    prediction = (prob_map > threshold).float()
+    prediction = (prob_map > 0.5).float()
     prob_map, prediction = prob_map.numpy(), prediction.numpy()
     save_image(image, 'image.png')
     save_image(prob_map, 'prob.png', cmap='jet')
@@ -41,14 +38,15 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(type='pil', label='TIFF Image')
-            threshold_slider = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.5, label='Prediction Threshold')
-            segment_button = gr.Button('Segment')
+            # threshold_slider = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.5, label='Prediction Threshold')
+            segment_button = gr.Button('Get Sidewalk Mask')
         with gr.Column():
-            prediction = gr.Image(type='pil', label='Predicted Mask')
+            mask = gr.Image(type='pil', label='Predicted Mask')
             prob_map = gr.Image(type='pil', label='Predicted Probability Map')
     segment_button.click(
         segment_sidewalk, 
-        inputs=[image_input, threshold_slider], 
-        outputs=[image_input, prediction, prob_map]
+        # inputs=[image_input, threshold_slider], 
+        inputs=[image_input], 
+        outputs=[image_input, mask, prob_map]
     )
 demo.launch(debug=True, show_error=True)