AntonMu
81 linhas
2.4 KiB
Python
81 linhas
2.4 KiB
Python
from PIL import Image
|
|
from os import path, makedirs
|
|
import os
|
|
import re
|
|
import pandas as pd
|
|
import sys
|
|
import argparse
|
|
|
|
|
|
def get_parent_dir(n=1):
|
|
""" returns the n-th parent dicrectory of the current
|
|
working directory """
|
|
current_path = os.path.dirname(os.path.abspath(__file__))
|
|
for k in range(n):
|
|
current_path = os.path.dirname(current_path)
|
|
return current_path
|
|
|
|
|
|
sys.path.append(os.path.join(get_parent_dir(1), "Utils"))
|
|
from Convert_Format import convert_vott_csv_to_yolo
|
|
|
|
Data_Folder = os.path.join(get_parent_dir(1), "Data")
|
|
VoTT_Folder = os.path.join(
|
|
Data_Folder, "Source_Images", "Training_Images", "vott-csv-export"
|
|
)
|
|
VoTT_csv = os.path.join(VoTT_Folder, "Annotations-export.csv")
|
|
YOLO_filename = os.path.join(VoTT_Folder, "data_train.txt")
|
|
|
|
model_folder = os.path.join(Data_Folder, "Model_Weights")
|
|
classes_filename = os.path.join(model_folder, "data_classes.txt")
|
|
|
|
if __name__ == "__main__":
|
|
# surpress any inhereted default values
|
|
parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS)
|
|
"""
|
|
Command line options
|
|
"""
|
|
parser.add_argument(
|
|
"--VoTT_Folder",
|
|
type=str,
|
|
default=VoTT_Folder,
|
|
help="Absolute path to the exported files from the image tagging step with VoTT. Default is "
|
|
+ VoTT_Folder,
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--VoTT_csv",
|
|
type=str,
|
|
default=VoTT_csv,
|
|
help="Absolute path to the *.csv file exported from VoTT. Default is "
|
|
+ VoTT_csv,
|
|
)
|
|
parser.add_argument(
|
|
"--YOLO_filename",
|
|
type=str,
|
|
default=YOLO_filename,
|
|
help="Absolute path to the file where the annotations in YOLO format should be saved. Default is "
|
|
+ YOLO_filename,
|
|
)
|
|
|
|
FLAGS = parser.parse_args()
|
|
|
|
# Prepare the dataset for YOLO
|
|
multi_df = pd.read_csv(FLAGS.VoTT_csv)
|
|
labels = multi_df["label"].unique()
|
|
labeldict = dict(zip(labels, range(len(labels))))
|
|
multi_df.drop_duplicates(subset=None, keep="first", inplace=True)
|
|
train_path = FLAGS.VoTT_Folder
|
|
convert_vott_csv_to_yolo(
|
|
multi_df, labeldict, path=train_path, target_name=FLAGS.YOLO_filename
|
|
)
|
|
|
|
# Make classes file
|
|
file = open(classes_filename, "w")
|
|
|
|
# Sort Dict by Values
|
|
SortedLabelDict = sorted(labeldict.items(), key=lambda x: x[1])
|
|
for elem in SortedLabelDict:
|
|
file.write(elem[0] + "\n")
|
|
file.close()
|