from AnnotationFormatConverter import vocxml_to_yolotxt

# define class name to id dict
class_name_to_id_mapping = {"smoke": 0,
                            "fire": 1}

# get the annotations
annotation_dir = 'annotations_xml/'
saved_dir = 'annotations_txt/val/'
annotations = glob.glob(annotation_dir + '*.xml')
annotations.sort()

# convert and save the annotations
for ann in tqdm.tqdm(annotations):
    info_dict = vocxml_to_yolotxt.extract_info_from_xml(ann)
    vocxml_to_yolotxt.convert_to_yolo_txt(saved_dir, info_dict, class_name_to_id_mapping)

from AnnotationFormatConverter import yolotxt_to_vocxml

# define id to class mapping dict
id_to_class_name_mapping = {0: "smoke",
                            1: "fire"}

# define dataset directory that contain images and labels in the same folder
data_dir = 'data_dir'
saved_dir = 'annotation_xml'

yolotxt_to_vocxml.convert_to_voc_xml(id_to_class_name_mapping, data_dir, 'saved_dir')

from ImageClustering import image_similarity

# get list of images
images_dir = "data_neuroo/train/images/"
images_path = sorted(os.listdir(images_dir))

# do image clustering
image_clusters, image_centroids, images_duplicate = image_similarity.cluster(root=images_dir, images_list=images_path,, sim_threshold=0.93,  
                                                    min_community_size=1, emb_batch_size=64, cluster_batch_size=128, size=320)

from DatasetStatistic import dataset_stistics

# define class names
names = {0: '0',
        1: '1', 
        2: '2',
        3: '3'}

# get list of labels
labels_dir = "train/labels/"
labels_path = sorted(os.listdir(labels_dir))

labels = []

# read labels
for i, label_file in enumerate(labels_path):
    with open(labels_dir + label_file, 'r') as fl:
        lines = fl.readlines()
    fl.close()

    for line in lines:
        line = line.split(" ")
        line = [float(i) for i in line]
        labels.append(line)
    
labels = np.array(labels)

# plot labels statistic
plot_labels_statistic(labels, names)

from DatasetSplit import approximation_multi_label_data_split

# list label files
labels_dir = "train/labels/"
labels_path = sorted(os.listdir(labels_dir))

# split train/test
train_index, test_index = approximation_multi_label_data_split.split(labels_path=labels_path, labels_dir=labels_dir, n_splits=7, shuffle=True, random_state=0)