Update app.py

app.py

@@ -9,9 +9,7 @@ from samgeo import SamGeo, show_image, download_file, overlay_images, tms_to_geo
 from zipfile import ZipFile
 
 
-'''
-Here we are using a pretrained model
-'''
+#initialize the model
 sam = SamGeo(
     model_type="vit_h",
     checkpoint='sam_vit_h_4b8939.pth',
@@ -35,10 +33,10 @@ def get_all_file_paths(directory):
 
 def get_shape_files(image):
   filename, file_extension = os.path.splitext(image.name)
-  print(filename,file_extension)
-  if file_extension=='.tif':
+  # print(filename,file_extension)
+  if file_extension in ['.tif','tiff']:
     #mkdir for app
-    storage='app_store'
+    storage='ShapeFiles_Folder'
     if os.path.exists(storage):
         pass
     else:
@@ -59,6 +57,11 @@ def get_shape_files(image):
     # load_model()
     sam.generate(image.name, output=f"{out_file}/masks.tif", foreground=True, unique=True)
     sam.tiff_to_vector(f"{out_file}/masks.tif", out_file)
+    sam.show_anns(axis="off", alpha=1, output=f"{out_file}/annotations.tif")
+
+    #converting tif to jpg
+    image = cv2.imread(f"{out_file}/annotations.tif")
+    cv2.imwrite("{}/{}.jpg".format(out_file,'converted'), image)
 
     # writing files to a zipfile
     file_paths=get_all_file_paths(out_file)
@@ -73,7 +76,7 @@ def get_shape_files(image):
 
     print('All files zipped successfully!')
 
-    return f'{zip_folder}/shapefile.zip'
+    return f"{out_file}/converted.jpg",f'{zip_folder}/shapefile.zip'
   # get_shape_files(image)
   else:
     return "Try uploading .tif file for processing!."
@@ -81,13 +84,14 @@ def get_shape_files(image):
 my_app = gr.Blocks()
 with my_app:
   gr.Markdown("<center><h1>Farmland Segmentation using Satellite Image data</h1></center>")
-  gr.Markdown("<center><h3>Provided with a .tif image file of satellite info, this app can segment the farmlands</h></center>")
+  gr.Markdown("<center><h3>Provided with a .tif image file of satellite, this app can segment the farmlands</h></center>")
   with gr.TabItem("Get Shapefiles"):
     with gr.Row():
       with gr.Column():
           img_source = gr.File(label="Please select source tif")
           source_image_loader = gr.Button("Get Shape File")
       with gr.Column():
+          image=gr.Image()
           output = gr.outputs.File("zip")
-  source_image_loader.click(get_shape_files,img_source,output)
+  source_image_loader.click(get_shape_files,img_source,[image,output])
 my_app.launch(debug = True)