File size: 2,433 Bytes
ea60552
 
 
 
 
 
 
 
 
 
c31846b
ea60552
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c31846b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ea60552
 
 
 
 
c31846b
 
 
ea60552
 
 
 
 
 
 
c31846b
ea60552
 
 
 
 
 
c31846b
ea60552
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
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,
)

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)
  print(filename,file_extension)
  if file_extension=='.tif':
    #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)
  
    out_file=f'{storage}/out_files'
    if os.path.exists(out_file):
        pass
    else:
      os.mkdir(out_file)

    # 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)

    # 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}/shapefile.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}/shapefile.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)