app.py

import gradio as gr
import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2
import os
import leafmap
from samgeo import SamGeo, show_image, download_file, overlay_images, tms_to_geotiff
from zipfile import ZipFile

#initialize the model
sam = SamGeo(
    model_type="vit_h",
    checkpoint='sam_vit_h_4b8939.pth',
    sam_kwargs=None,
)

#mkdir for app
storage='app_store'
if os.path.exists(storage):
    pass
else:
  os.mkdir(storage)

zip_folder=f'{storage}/zip_folder'
if os.path.exists(zip_folder):
    pass
else:
  os.mkdir(zip_folder)

def get_all_file_paths(directory):
  
    # initializing empty file paths list
    file_paths = []
  
    # crawling through directory and subdirectories
    for root, directories, files in os.walk(directory):
        for filename in files:
            # join the two strings in order to form the full filepath.
            filepath = os.path.join(root, filename)
            file_paths.append(filepath)
  
    # returning all file paths
    return file_paths  

def get_shape_files(image):
  filename, file_extension = os.path.splitext(image.name)
  if file_extension=='tif':
    out_file=f'{storage}/{image.name}'
    if os.path.exists(out_file):
        pass
    else:
      os.mkdir(out_file)

    sam.generate(image.name, output="masks.tif", foreground=True, unique=True)
    sam.tiff_to_vector("masks.tif", out_file)

    # writing files to a zipfile
    file_paths=get_all_file_paths(out_file)
    
    # for file_name in file_paths:
    #       print(file_name)

    with ZipFile(f'{zip_folder}/{image.name}.zip','w') as zip:
        # writing each file one by one
        for file in file_paths:
            zip.write(file)
    
    print('All files zipped successfully!') 

    return f'{zip_folder}/{image.name}.zip'
  # get_shape_files(image)
  else:
    return "Try uploading .tif file for processing!."

my_app = gr.Blocks()
with my_app:
  gr.Markdown("Segmenting Satellite Image")
  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():
          output = gr.outputs.File("zip")
  source_image_loader.click(get_shape_files,img_source,output)
my_app.launch(debug = True)