#!/usr/bin/python3
# -*- encoding: utf-8 -*-
"""
Segments images using a pre-trained ONNX network;
the network is trained to segment aerial images into 9 classes;
see https://github.com/eokeeffe/UAV_Aerial_Segmentation_cpp_onnx

Install:
  pip install opencv-python-headless onnxruntime numpy tqdm argparse pathlib

Example usage:
  python3 ImageSegmentation.py -i images -o masks

In order to use the segmentation masks to segment the dense point-cloud, add these extra params:
  DensifyPointCloud scene.mvs  -m masks --estimate-segmentation 2 -v 3

Created by @eokeeffe
"""

import argparse
import cv2
import json
import numpy as np
import os
import onnxruntime as ort
from pathlib import Path
from tqdm import tqdm

def loadImage(image_name):
  image = cv2.imread(image_name, cv2.IMREAD_UNCHANGED)
  height,width = image.shape[:2]

  # image dims have to be 1024,576
  my_image_test = cv2.resize(image, (1024,576), interpolation=cv2.INTER_LINEAR)
  # need to be floating point
  my_image_test = my_image_test.astype('float32')
  my_image_test /= 255.0
  # apply the normalization from pytorch
  mean=[0.485, 0.456, 0.406]
  std=[0.229, 0.224, 0.225]

  my_image_test[..., 0] -= mean[0]
  my_image_test[..., 1] -= mean[1]
  my_image_test[..., 2] -= mean[2]

  my_image_test[..., 0] /= std[0]
  my_image_test[..., 1] /= std[1]
  my_image_test[..., 2] /= std[2]
  my_image_test = my_image_test.transpose(2, 0, 1)
  my_image_test = np.expand_dims(my_image_test, axis=0)
  # final dims should be 1,3,576,1024
  return my_image_test,height,width

def extractSegmentedImage(outputs, original_height, original_width, sigmoid_threshold = 0.8):
  output_masks = outputs[0].transpose(1, 2, 0)
  segmented_image = np.zeros((original_height, original_width),dtype=np.uint8)
  for ch in range(output_masks.shape[-1]):
    seg_mask = output_masks[:,:,ch]
    seg_mask[seg_masksigmoid_threshold] = 1
    seg_mask = seg_mask.astype(np.uint8)
    seg_mask = cv2.resize(seg_mask, (original_width,original_height), interpolation= cv2.INTER_LINEAR)
    indxs = np.where(seg_mask>0)
    segmented_image[indxs] = ch+1
  return segmented_image

def createPxielLabels():
  label_json = {
    "0": "unclassified",
    "1": "clutter",
    "2": "building",
    "3": "road",
    "4": "static_car",
    "5": "tree",
    "6": "vegetation",
    "7": "human",
    "8": "moving_car"
  }
  return label_json

def segmentImages(images_path, output_path, onnx_file, labels_file, sigmoid_threshold=0.8):
  # check if the onnx network exists
  if(not os.path.exists(onnx_file)):
    # download the onnx network
    import urllib.request
    url = "https://github.com/eokeeffe/UAV_Aerial_Segmentation_cpp_onnx/raw/refs/heads/main/networks/aerial_segmentation.onnx"
    if not os.path.isabs(onnx_file):
      onnx_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), onnx_file)
    print(f"Downloading segmentation model to {onnx_file}...")
    urllib.request.urlretrieve(url, onnx_file)

  # load the onnx network
  ort_session = ort.InferenceSession(onnx_file)

  # get the image locations
  all_images = os.listdir(images_path)

  # create the output folder if it doesn't exist
  Path(output_path).mkdir(parents=True, exist_ok=True)

  # segment each image
  print("Starting segmentation ...")
  for image in tqdm(all_images):
    input_image = os.path.join(images_path, image)
    output_image = os.path.join(output_path, os.path.splitext(image)[0] + '.mask.png')
    
    if(not os.path.exists(input_image)):
      print(input_image," doesn't exist")
      continue
    if(os.path.exists(output_image)):
      print(output_image," already exists")
      continue

    # format the image to the correct dimensions
    preprocessed_image,h,w = loadImage(input_image)
    # run the inference
    outputs = ort_session.run(["sigmoid"], {'image': preprocessed_image})[0]
    # process the output to classified pixels
    classified_image = extractSegmentedImage(outputs, h, w, sigmoid_threshold=sigmoid_threshold)
    # save the segmented image
    cv2.imwrite(output_image, classified_image)

  # save a json file with the pixel value to label relationship
  if labels_file is not None:
    if not os.path.isabs(labels_file):
      labels_file = os.path.join(output_path, labels_file)
    with open(labels_file, "w") as outfile:
      json.dump(createPxielLabels(), outfile)

  ort_session = None
  print("... segmentation completed!")
  
if __name__=="__main__":
  parser = argparse.ArgumentParser()
  parser.add_argument("-i", "--images", help = "directory with images to do semantic segmentations")
  parser.add_argument("-o", "--output", help = "directory to store the segmented images")
  parser.add_argument("-n", "--onnx", default='aerial_segmentation.onnx', help = "onnx network to use")
  parser.add_argument("-l", "--labels", default='labels.json', help = "export label names to json file")
  parser.add_argument("-s", "--sigmoid", default=0.8, help = "sigmoid threshold")
  args = parser.parse_args()

  segmentImages(args.images, args.output, args.onnx, args.labels, float(args.sigmoid))