Update app.py

app.py

@@ -2,14 +2,10 @@ import os
 import cv2
 from PIL import Image
 import numpy as np
-from keras import backend as K
-from matplotlib import pyplot as plt
-from tensorflow.keras.utils import to_categorical
 import geopandas as gpd
 import matplotlib.pyplot as plt
 from keras.models import load_model
 from tensorflow.keras.preprocessing.image import load_img, img_to_array
-from google.colab.patches import cv2_imshow
 import geopandas as gpd
 from skimage.measure import regionprops, label
 from shapely.geometry import Polygon
@@ -29,12 +25,7 @@ def predict(img):
     return predictions,predicted_image
 
 def get_shape_files(img):
-  # predictions,_ = predict(img)
-  model = load_model('building_footprint_extraction_model.h5')
-  img_array = img_to_array(img)
-  img_array = img_array.reshape((1, 256, 256, 3))
-  img_array = img_array / 255.0
-  predictions = model.predict(img_array)
+  predictions,predicted_image= predict(img)
   threshold = 0.5
   binary_mask = (predictions > threshold).astype(np.uint8)[:, :, 1]
   if np.sum(binary_mask) == 0:
@@ -55,9 +46,6 @@ def get_shape_files(img):
       gdf.to_file(output_shapefile)
 
       # To get Masked Image
-      predicted_image = np.argmax(predictions, axis=3)
-      predicted_image = predicted_image[0,:,:]
-      predicted_image = predicted_image * 255
       cv2.imwrite('shapefiles/mask.jpg',predicted_image)
       shutil.make_archive('shapefile', 'zip', 'shapefiles')
       return 'shapefile.zip',predicted_image
@@ -66,25 +54,25 @@ my_app = gr.Blocks()
 with my_app:
   gr.Markdown("<center><h1>Building Footprint Extraction</h1></center>")
   with gr.Tabs():
-    with gr.TabItem("Get Mask Image"):
-      with gr.Row():
-        with gr.Column():
-            img_source = gr.Image(label="Please select source Image", shape=(256, 256))
-            source_image_loader = gr.Button("Load above Image")
-        with gr.Column():
-            img_output = gr.Image(label="Image Output")
-    source_image_loader.click(predict,img_source,img_output)
-    with gr.TabItem("Get Shapefiles"):
+    # with gr.TabItem("Get Mask Image"):
+    #   with gr.Row():
+    #     with gr.Column():
+    #         img_source = gr.Image(label="Please select source Image", shape=(256, 256))
+    #         source_image_loader = gr.Button("Load above Image")
+    #     with gr.Column():
+    #         img_output = gr.Image(label="Image Output")
+    # source_image_loader.click(predict,img_source,[img_output])
+    with gr.TabItem("Get Mask Image and Shapefiles"):
       with gr.Row():
         with gr.Column():
             img_source = gr.Image(label="Please select source Image", shape=(256, 256))
             get_shape_loader = gr.Button("Get Shape File")
-        
+
         with gr.Column():
           with gr.Row():
               mask_img=gr.Image(label="Image Output")
           with gr.Row():
             output_zip = gr.outputs.File()
-            
+
     get_shape_loader.click(get_shape_files,img_source,[output_zip,mask_img])
 my_app.launch(debug = True)