Update app.py

app.py

@@ -4,29 +4,15 @@ import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
+import torchvision
 import glob
 import gradio as gr
 from PIL import Image
-
 from segment_anything import SamAutomaticMaskGenerator, SamPredictor, sam_model_registry
-from torchvision.utils import draw_segmentation_masks
 
 matplotlib.pyplot.switch_backend('Agg') # for matplotlib to work in gradio
-
-gr.Markdown(
-    """
-# Segment Anything Model (SAM) 
-### A test on remote sensing data
-
-- Paper:[Link](https://scontent-fml2-1.xx.fbcdn.net/v/t39.2365-6/10000000_900554171201033_1602411987825904100_n.pdf?_nc_cat=100&ccb=1-7&_nc_sid=3c67a6&_nc_ohc=Ald4OYhL6hgAX-ZcGmS&_nc_ht=scontent-fml2-1.xx&oh=00_AfDk4FvyiDYeXgflANA2CbdV6HSS8CcJmrvjSfTqsgUmog&oe=643500E7)
-- Github:[https://github.com/facebookresearch/segment-anything](https://github.com/facebookresearch/segment-anything)
-- Dataset:https://ai.facebook.com/datasets/segment-anything-downloads/(https://ai.facebook.com/datasets/segment-anything-downloads/)
-- Official Demo:[https://segment-anything.com/demo](https://segment-anything.com/demo)
-"""
-)
-
 #setup model
-sam_checkpoint = "meta-model.pth"
+sam_checkpoint = "sam_vit_h_4b8939.pth"
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # use GPU if available
 model_type = "default"
 sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
@@ -62,17 +48,24 @@ def segment_image(image):
     plt.savefig('output.png', bbox_inches='tight', pad_inches=0)
     output = cv2.imread('output.png')
     return Image.fromarray(output)
-
+    
 with gr.Blocks() as demo:
-    # gr.Markdown("## Segment-anything Demo")
-
+    gr.Markdown(
+        """
+    # Segment Anything Model (SAM) 
+    ### A test on remote sensing data
+    - Paper:[(https://arxiv.org/abs/2304.02643](https://arxiv.org/abs/2304.02643)
+    - Github:[https://github.com/facebookresearch/segment-anything](https://github.com/facebookresearch/segment-anything)
+    - Dataset:https://ai.facebook.com/datasets/segment-anything-downloads/(https://ai.facebook.com/datasets/segment-anything-downloads/)
+    - Official Demo:[https://segment-anything.com/demo](https://segment-anything.com/demo)
+    """
+    )
     with gr.Row():
         image = gr.Image()
         image_output = gr.Image()
-
+    print(image.shape)
     segment_image_button = gr.Button("Segment")
     segment_image_button.click(segment_image, inputs=[image], outputs=image_output)
     gr.Examples(glob.glob('./images/*'),image,image_output,segment_image)
-
     
-demo.launch()
+demo.launch(debug=True)