vishnu23 commited on
Commit
ea60552
·
1 Parent(s): bddd29e

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +87 -0
app.py ADDED
@@ -0,0 +1,87 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import numpy as np
3
+ import torch
4
+ import matplotlib.pyplot as plt
5
+ import cv2
6
+ import os
7
+ import leafmap
8
+ from samgeo import SamGeo, show_image, download_file, overlay_images, tms_to_geotiff
9
+ from zipfile import ZipFile
10
+
11
+ #initialize the model
12
+ sam = SamGeo(
13
+ model_type="vit_h",
14
+ checkpoint='sam_vit_h_4b8939.pth',
15
+ sam_kwargs=None,
16
+ )
17
+
18
+ #mkdir for app
19
+ storage='app_store'
20
+ if os.path.exists(storage):
21
+ pass
22
+ else:
23
+ os.mkdir(storage)
24
+
25
+ zip_folder=f'{storage}/zip_folder'
26
+ if os.path.exists(zip_folder):
27
+ pass
28
+ else:
29
+ os.mkdir(zip_folder)
30
+
31
+ def get_all_file_paths(directory):
32
+
33
+ # initializing empty file paths list
34
+ file_paths = []
35
+
36
+ # crawling through directory and subdirectories
37
+ for root, directories, files in os.walk(directory):
38
+ for filename in files:
39
+ # join the two strings in order to form the full filepath.
40
+ filepath = os.path.join(root, filename)
41
+ file_paths.append(filepath)
42
+
43
+ # returning all file paths
44
+ return file_paths
45
+
46
+ def get_shape_files(image):
47
+ filename, file_extension = os.path.splitext(image.name)
48
+ if file_extension=='tif':
49
+ out_file=f'{storage}/{image.name}'
50
+ if os.path.exists(out_file):
51
+ pass
52
+ else:
53
+ os.mkdir(out_file)
54
+
55
+ sam.generate(image.name, output="masks.tif", foreground=True, unique=True)
56
+ sam.tiff_to_vector("masks.tif", out_file)
57
+
58
+ # writing files to a zipfile
59
+ file_paths=get_all_file_paths(out_file)
60
+
61
+ # for file_name in file_paths:
62
+ # print(file_name)
63
+
64
+ with ZipFile(f'{zip_folder}/{image.name}.zip','w') as zip:
65
+ # writing each file one by one
66
+ for file in file_paths:
67
+ zip.write(file)
68
+
69
+ print('All files zipped successfully!')
70
+
71
+ return f'{zip_folder}/{image.name}.zip'
72
+ # get_shape_files(image)
73
+ else:
74
+ return "Try uploading .tif file for processing!."
75
+
76
+ my_app = gr.Blocks()
77
+ with my_app:
78
+ gr.Markdown("Segmenting Satellite Image")
79
+ with gr.TabItem("Get Shapefiles"):
80
+ with gr.Row():
81
+ with gr.Column():
82
+ img_source = gr.File(label="Please select source tif")
83
+ source_image_loader = gr.Button("Get Shape File")
84
+ with gr.Column():
85
+ output = gr.outputs.File("zip")
86
+ source_image_loader.click(get_shape_files,img_source,output)
87
+ my_app.launch(debug = True)