# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # DeepWay prediction and Path Planning

# %matplotlib inline
# %reload_ext autoreload
# %autoreload 2

# import some data
import os
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import tensorflow as tf
from glob import glob
import cv2
from tqdm.notebook import tqdm

from utils.geometry import *
from utils.visualization import visualize_mask
from utils.postprocessing import wp_post_processing
from utils.planning import AStarPlanner,plan_path,field_coverage
from utils.tools import load_config,read_and_predict
from utils.deep_way_net import build_deepway

gpus = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_visible_devices(gpus[0], 'GPU')
tf.config.experimental.set_memory_growth(gpus[0], True)

# +
# important paths and names
PATH_DIR = os.path.abspath('.')

TEST_DATA_PATH = os.path.join(PATH_DIR, 'dataset/test')
SATELLITE_DATA_PATH = os.path.join(PATH_DIR, 'dataset/satellite')
PATH_WEIGHTS = os.path.join(PATH_DIR, 'bin')

#name_model = 'deep_way_pretrained.h5'
name_model = 'deep_way.h5'


waypoint_file_name = 'waypoints_gt.csv'
config_file = 'utils/config.json'
# -

# important config
config = load_config(config_file)

# +
# other important variables

#Deepway
K = config['K']
CONF_THR = 0.9
DIST_THR = 8

#Astar
grid_size=1
robot_radius=1.0
W=H=config['MASK_DIM']
# -

# # Load DeepWay

# load network
deepway_net = build_deepway(name_model, config['FILTERS'], config['KERNEL_SIZE'],
                            config['N'], config['MASK_DIM'])
deepway_net.load_weights(os.path.join(PATH_WEIGHTS, name_model))

# # Path planning on a single image

# +
img = 0

img_folder = SATELLITE_DATA_PATH  #for satellite test set
#img_folder = TEST_DATA_PATH       #for artificial test set

img_list = glob(img_folder+"/img*.png")
img_list = sorted(img_list,key=lambda s: int(s.split('/')[-1][3:-4])) #sort by name
img = img_list[img]
img
# -

wp_gt_df = pd.read_csv(img_folder+"/waypoints.csv")
name = img.split('/')[-1][:-4]
x = wp_gt_df.loc[wp_gt_df['N_img'] == name]["x_wp"].to_numpy()[:,None]
y = wp_gt_df.loc[wp_gt_df['N_img'] == name]["y_wp"].to_numpy()[:,None]
wp_gt = np.concatenate((x,y),axis=-1)
img,wp = read_and_predict(img,deepway_net,CONF_THR,DIST_THR,K)
order,cluster_order = wp_post_processing(img,wp)

visualize_mask(img,path=order)
#visualize_mask(img,wp=ground_truth)

mask = np.bitwise_not(img.astype("bool")) #get mask with obstacles as ones 
obs_y,obs_x = np.nonzero(mask) #get obstacles coordinates
a_star = AStarPlanner(obs_x,obs_y,grid_size,robot_radius,H,W)
path = plan_path(order,a_star,smooth=False)

cov = field_coverage(mask,path,order,cluster_order,wp_gt)
print("Coverage:",cov)
visualize_mask(img,wp=path,rad=2)

# # Path planning on the entire datasets and coverage metric

# +
img_folder = SATELLITE_DATA_PATH  #for satellite test set
#img_folder = TEST_DATA_PATH       #for artificial test set

img_list = glob(img_folder+"/img*.png")
img_list = sorted(img_list,key=lambda s: int(s.split('/')[-1][3:-4])) #sort by name
#img_list
# -

# ## Prediction and post-processing

# +
images = []
ground_truths = []
waypoints = []
orders = []
cluster_orders = []

wp_gt = pd.read_csv(img_folder+"/waypoints.csv")

for img in tqdm(img_list):
    name = img.split('/')[-1][:-4]
    x = wp_gt.loc[wp_gt['N_img'] == name]["x_wp"].to_numpy()[:,None]
    y = wp_gt.loc[wp_gt['N_img'] == name]["y_wp"].to_numpy()[:,None]
    ground_truths.append(np.concatenate((x,y),axis=-1))
    img,wp = read_and_predict(img,deepway_net,CONF_THR,DIST_THR,K)
    images.append(img)
    waypoints.append(wp)
    p,c = wp_post_processing(img,wp)
    orders.append(p)
    cluster_orders.append(c)
# -

index = 0
visualize_mask(images[index],wp=waypoints[index])
#visualize_mask(images[index],wp=ground_truths[index])

visualize_mask(images[index],path=orders[index])

# ## Path planning (A*)

paths = []
for i in tqdm(range(len(images))):
    img = images[i]
    p = orders[i]
    mask = np.bitwise_not(img.astype("bool"))
    obs_y,obs_x = np.nonzero(mask)
    a_star = AStarPlanner(obs_x,obs_y,grid_size,robot_radius,H,W)
    paths.append(plan_path(p,a_star,img_index=i,smooth=False))

visualize_mask(images[index],wp=paths[index],rad=2)

# ## Coverage metric 

coverages = []
for img,p,wp,c,wp_gt in tqdm(zip(images,paths,orders,cluster_orders,ground_truths),total=len(images)):
    mask = np.bitwise_not(img.astype("bool"))
    coverages.append(field_coverage(mask,p,wp,c,wp_gt))
coverages = np.array(coverages)

print("Mean coverage:",np.mean(coverages))
plt.hist(coverages,20)
plt.show()

StarcoderdataJupyterScriptsDedupFiltered