import argparse

import os

import yaml

import sys

import numpy as np

DISTANCES = [(1e-8, 10.0),

(40.0, 50.0)]

splits = ["train", "valid", "test"]

backends = ["numpy", "torch"]

if __name__ == '__main__':

parser = argparse.ArgumentParser("./evaluate_semantics_by_distance.py")

parser.add_argument(

'--dataset', '-d',

type=str,

required=True,

help='Dataset dir. No Default',

)

parser.add_argument(

'--predictions', '-p',

type=str,

required=None,

help='Prediction dir. Same organization as dataset, but predictions in'

'each sequences "prediction" directory. No Default. If no option is set'

' we look for the labels in the same directory as dataset'

)

parser.add_argument(

'--split', '-s',

type=str,

required=False,

default="valid",

help='Split to evaluate on. One of ' +

str(splits) + '. Defaults to %(default)s',

)

parser.add_argument(

'--backend', '-b',

type=str,

required=False,

default="numpy",

help='Backend for evaluation. One of ' +

str(backends) + ' Defaults to %(default)s',

)

parser.add_argument(

'--datacfg', '-dc',

type=str,

required=False,

default="config/semantic-kitti.yaml",

help='Dataset config file. Defaults to %(default)s',

)

parser.add_argument(

'--limit', '-l',

type=int,

required=False,

default=None,

help='Limit to the first "--limit" points of each scan. Useful for'

' evaluating single scan from agregated pointcloud.'

' Defaults to %(default)s',

)

parser.add_argument(

'--codalab',

dest='codalab',

default=False,

action='store_true',

help='Exports "segmentation_scores_distance.txt" for codalab'

'Defaults to %(default)s',

)

FLAGS, unparsed = parser.parse_known_args()

# fill in real predictions dir

if FLAGS.predictions is None:

FLAGS.predictions = FLAGS.dataset

# print summary of what we will do

print("*" * 80)

print("INTERFACE:")

print("Data: ", FLAGS.dataset)

print("Predictions: ", FLAGS.predictions)

print("Backend: ", FLAGS.backend)

print("Split: ", FLAGS.split)

print("Config: ", FLAGS.datacfg)

print("Limit: ", FLAGS.limit)

print("Codalab: ", FLAGS.codalab)

print("*" * 80)

# assert split

assert(FLAGS.split in splits)

# assert backend

assert(FLAGS.backend in backends)

# open data config file

try:

print("Opening data config file %s" % FLAGS.datacfg)

DATA = yaml.safe_load(open(FLAGS.datacfg, 'r'))

except Exception as e:

print(e)

print("Error opening data yaml file.")

quit()

# get number of interest classes, and the label mappings

class_strings = DATA["labels"]

class_remap = DATA["learning_map"]

class_inv_remap = DATA["learning_map_inv"]

class_ignore = DATA["learning_ignore"]

nr_classes = len(class_inv_remap)

# make lookup table for mapping

maxkey = 0

for key, data in class_remap.items():

if key > maxkey:

maxkey = key

# +100 hack making lut bigger just in case there are unknown labels

remap_lut = np.zeros((maxkey + 100), dtype=np.int32)

for key, data in class_remap.items():

try:

remap_lut[key] = data

except IndexError:

print("Wrong key ", key)

# print(remap_lut)

# create evaluator

ignore = []

for cl, ign in class_ignore.items():

if ign:

x_cl = int(cl)

ignore.append(x_cl)

print("Ignoring xentropy class ", x_cl, " in IoU evaluation")

# create evaluator

evaluators = []

for i in range(len(DISTANCES)):

if FLAGS.backend == "torch":

from auxiliary.torch_ioueval import iouEval

evaluators.append(iouEval(nr_classes, ignore))

evaluators[i].reset()

elif FLAGS.backend == "numpy":

from auxiliary.np_ioueval import iouEval

evaluators.append(iouEval(nr_classes, ignore))

evaluators[i].reset()

else:

print("Backend for evaluator should be one of ", str(backends))

quit()

# get test set

test_sequences = DATA["split"][FLAGS.split]

# get scan paths

scan_names = []

for sequence in test_sequences:

sequence = '{0:02d}'.format(int(sequence))

label_paths = os.path.join(FLAGS.dataset, "sequences",

str(sequence), "velodyne")

# populate the label names

seq_scan_names = [os.path.join(dp, f) for dp, dn, fn in os.walk(

os.path.expanduser(label_paths)) for f in fn if ".bin" in f]

seq_scan_names.sort()

scan_names.extend(seq_scan_names)

# print(scan_names)

# get label paths

label_names = []

for sequence in test_sequences:

sequence = '{0:02d}'.format(int(sequence))

label_paths = os.path.join(FLAGS.dataset, "sequences",

str(sequence), "labels")

# populate the label names

seq_label_names = [os.path.join(dp, f) for dp, dn, fn in os.walk(

os.path.expanduser(label_paths)) for f in fn if ".label" in f]

seq_label_names.sort()

label_names.extend(seq_label_names)

# print(label_names)

# get predictions paths

pred_names = []

for sequence in test_sequences:

sequence = '{0:02d}'.format(int(sequence))

pred_paths = os.path.join(FLAGS.predictions, "sequences",

sequence, "predictions")

# populate the label names

seq_pred_names = [os.path.join(dp, f) for dp, dn, fn in os.walk(

os.path.expanduser(pred_paths)) for f in fn if ".label" in f]

seq_pred_names.sort()

pred_names.extend(seq_pred_names)

# print(pred_names)

# check that I have the same number of files

print("scans", len(scan_names))

print("labels: ", len(label_names))

print("predictions: ", len(pred_names))

assert(len(label_names) == len(pred_names) and

len(scan_names) == len(label_names))

# open each file, get the tensor, and make the iou comparison

for scan_file, label_file, pred_file in zip(scan_names, label_names, pred_names):

print("evaluating scan ", scan_file)

# open scan

scan = np.fromfile(scan_file, dtype=np.float32)

scan = scan.reshape((-1, 4)) # reshape to matrix

if FLAGS.limit is not None:

scan = scan[:FLAGS.limit] # limit to desired length

depth = np.linalg.norm(scan, 2, axis=1) # get depth to filter by distance

# open label

label = np.fromfile(label_file, dtype=np.int32)

label = label.reshape((-1)) # reshape to vector

label = label & 0xFFFF # get lower half for semantics

if FLAGS.limit is not None:

label = label[:FLAGS.limit] # limit to desired length

label = remap_lut[label] # remap to xentropy format

# open prediction

pred = np.fromfile(pred_file, dtype=np.int32)

pred = pred.reshape((-1)) # reshape to vector

pred = pred & 0xFFFF # get lower half for semantics

if FLAGS.limit is not None:

pred = pred[:FLAGS.limit] # limit to desired length

pred = remap_lut[pred] # remap to xentropy format

# evaluate for all distances

for idx in range(len(DISTANCES)):

# select by range

lrange = DISTANCES[idx][0]

hrange = DISTANCES[idx][1]

mask = np.logical_and(depth > lrange, depth < hrange)

# mask by distance

# mask_depth = depth[mask]

# print("mask range, ", mask_depth.max(), mask_depth.min())

mask_label = label[mask]

mask_pred = pred[mask]

# add single scan to evaluation

evaluators[idx].addBatch(mask_pred, mask_label)

# print for all ranges

print("*" * 80)

for idx in range(len(DISTANCES)):

# when I am done, print the evaluation

m_accuracy = evaluators[idx].getacc()

m_jaccard, class_jaccard = evaluators[idx].getIoU()

# print for spreadsheet

sys.stdout.write('range {lrange}m to {hrange}m,'.format(lrange=DISTANCES[idx][0],

hrange=DISTANCES[idx][1]))

for i, jacc in enumerate(class_jaccard):

if i not in ignore:

sys.stdout.write('{jacc:.3f}'.format(jacc=jacc.item()))

sys.stdout.write(",")

sys.stdout.write('{jacc:.3f}'.format(jacc=m_jaccard.item()))

sys.stdout.write(",")

sys.stdout.write('{acc:.3f}'.format(acc=m_accuracy.item()))

sys.stdout.write('\n')

sys.stdout.flush()

# if codalab is necessary, then do it

if FLAGS.codalab:

results = {}

for idx in range(len(DISTANCES)):

# make string for distance

d_str = str(DISTANCES[idx][-1])+"m_"

# get values for this distance range

m_accuracy = evaluators[idx].getacc()

m_jaccard, class_jaccard = evaluators[idx].getIoU()

# put in dictionary

results[d_str+"accuracy_mean"] = float(m_accuracy)

results[d_str+"iou_mean"] = float(m_jaccard)

for i, jacc in enumerate(class_jaccard):

if i not in ignore:

results[d_str+"iou_"+class_strings[class_inv_remap[i]]] = float(jacc)

# save to file

with open('segmentation_scores_distance.txt', 'w') as yaml_file:

yaml.dump(results, yaml_file, default_flow_style=False)