updated app

app.py

@@ -8,23 +8,22 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 model_config = SamConfig.from_pretrained("facebook/sam-vit-base")
 processor = SamProcessor.from_pretrained('facebook/sam-vit-base')
-model = SamModel(config=model_config)
+# model = SamModel(config=model_config)
 model = SamModel.from_pretrained('kitooo/sidewalk-seg-base')
 model.to(device)
 
 def segment_sidewalk(image, threshold):
-    # init data
     width, height = image.size
     prompt = [0, 0, width, height]
     inputs = processor(image, input_boxes=[[prompt]], return_tensors='pt')
-    # make prediction
-    outputs = model(pixel_values=inputs['pixel_values'].to(device),
-                    input_boxes=inputs['input_boxes'].to(device),
-                    multimask_output=False)
+    # 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()
     prob_map, prediction = prob_map.numpy(), prediction.numpy()
-    # visualize results
     save_image(image, 'image.png')
     save_image(prob_map, 'prob.png', cmap='jet')
     save_image(prediction, 'mask.png', cmap='gray')
@@ -45,8 +44,8 @@ with gr.Blocks() as demo:
             threshold_slider = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.5, label='Prediction Threshold')
             segment_button = gr.Button('Segment')
         with gr.Column():
-            prediction = gr.Image(type='pil', label='Segmentation Result')
-            prob_map = gr.Image(type='pil', label='Probability Map')
+            prediction = 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],