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
import shutil
import gradio as gr

def predict(img):
    # model = load_model('drive/My Drive/building_footprint_extraction_model.h5')
    model = load_model('/content/drive/MyDrive/Colab Notebooks/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)
    predicted_image = np.argmax(predictions, axis=3)
    predicted_image = predicted_image[0,:,:]
    predicted_image = predicted_image * 255
    return predictions,predicted_image

def get_shape_files(img):
  # predictions,_ = predict(img)
  model = load_model('/content/drive/MyDrive/Colab Notebooks/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)
  threshold = 0.5
  binary_mask = (predictions > threshold).astype(np.uint8)[:, :, 1]
  if np.sum(binary_mask) == 0:
      print("No building pixels detected. Saving an empty shapefile.")
  else:
      labeled_mask = label(binary_mask)
      building_polygons = []
      props = regionprops(labeled_mask)
      for prop in props:
          polygon = Polygon([(point[1], point[0]) for point in prop.coords])
          building_polygons.append(polygon)
      gdf = gpd.GeoDataFrame(geometry=building_polygons, crs="EPSG:4326")
      output_shapefile = "shapefiles/building_footprints.shp"
      if os.path.exists('shapefiles'):
        pass
      else:
        os.mkdir('shapefiles')
      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

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