# -*- coding: utf-8 -*-
"""lab1_sky_detector.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1fbwiWdURCSaRB3fbSs0gF88hiFqpJhC1
"""

#!pip install gradio

import gradio as gr
#from google.colab import files
import cv2
import numpy as np
#from matplotlib import pyplot as plt
#from scipy.signal import medfilt

# upload input image (in Google Colab environment)
#def upload_image():
#  uploaded = files.upload()
#  file = next(iter(uploaded))
#  img = cv2.imread(file)
#  return img

# display input image (in Google Colab environment)
#def display_image(img, title="Image"):
#  img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # convert from BGR to RGB for OpenCV
#  plt.imshow(img_rgb)
#  plt.title(title)
#  plt.axis("off")
#  plt.show()

# convert input image to grayscale to prepare for edge detection
def convert_to_grayscale(img):
  img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  return img_gray

# use Gaussian blurring to reduce noise
def apply_blur(img_gray):
  img_blurred = cv2.blur(img_gray, (3, 3))
  return img_blurred

# use Laplacian gradient to detect edges
def calculate_gradient(img_blurred, threshold=9):
  laplacian = cv2.Laplacian(img_blurred, cv2.CV_8U)
  gradient_mask = (laplacian < threshold).astype(np.uint8)
  return gradient_mask

# use morphological erosion to refine the mask
def refine_mask_morphological(gradient_mask):
  kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
  mask = cv2.morphologyEx(gradient_mask, cv2.MORPH_ERODE, kernel)
#  skyline_mask = cal_skyline(eroded_mask)
  return mask

# apply the mask to isolate the sky
#def cal_skyline(mask):
#  h, w = mask.shape
#  for i in range(w):
#    column = mask[:, i]
#    after_median = medfilt(column, kernel_size=21)
#    try:
#      first_white_index = np.where(after_median == 1)[0][0]
#      first_black_index = np.where(after_median == 0)[0][0]
#      if first_black_index > first_white_index:
#        mask[:first_black_index, i] = 1
#        mask[first_black_index:, i] = 0
#    except IndexError:
#      continue
#  return mask

# use median blur to refine the mask
def refine_mask_median(mask):
  h, w = mask.shape
  mask_blurred = cv2.medianBlur(mask, 9)
  for i in range(w):
    median = mask_blurred[:, i]
    try:
      first_white_index = np.where(median == 0)[0][0]
      first_black_index = np.where(median == 1)[0][0]
      if first_white_index > 20:
        mask[first_black_index:first_white_index, i] = 1
        mask[first_white_index:, i] = 0
        mask[:first_black_index, i] = 0
    except:
      continue
  return mask

# extract sky region by applying the mask
def get_sky_region(img, mask):
  sky_region = cv2.bitwise_and(img, img, mask=mask)
  return sky_region

# run in order and show the detected sky region
#def sky_detector(img):
#  display_image(img, "Original Image")
#  img_blurred = convert_to_grayscale_and_blur(img)
#  gradient_mask = calculate_gradient(img_blurred)
#  skyline_mask = refine_skyline(gradient_mask)
#  sky_region = get_sky_region(img, skyline_mask)
#  display_image(sky_region, "Sky Region")

# main
#image = upload_image()
#sky_detector(image)

# run in order for Gardio
def sky_detector(img):
  img_gray = convert_to_grayscale(img)
  img_blurred = apply_blur(img_gray)
  gradient_mask = calculate_gradient(img_blurred)
  mask = refine_mask_morphological(gradient_mask)
  mask = refine_mask_median(mask)
  sky_region = get_sky_region(img, mask)
#  sky_region_rgb = cv2.cvtColor(sky_region, cv2.COLOR_BGR2RGB)
  return sky_region

# set up Gardio interface
interface = gr.Interface(
    fn=sky_detector,
    inputs="image",
    outputs="image",
    title="Sky Detector",
    live=True
    )

interface.launch(share=True, debug=True)