diff --git a/src/concept.py b/src/concept.py index d11ca83..42d240f 100644 --- a/src/concept.py +++ b/src/concept.py @@ -7,3 +7,4 @@ class Concept(Enum): STOP = auto() RECHTS_ABBIEGEN = auto() LINKS_ABBIEGEN = auto() + RECHTS_VOR_LINKS = auto() diff --git a/src/create_feature_vectors.py b/src/create_feature_vectors.py new file mode 100644 index 0000000..4e99cb0 --- /dev/null +++ b/src/create_feature_vectors.py @@ -0,0 +1,61 @@ +import os +import cv2 as cv + +from featue_extraction import get_color_percentage, get_raster_color_percentage +from feature_vector import FeatureVector +from concept import Concept +from feature import Feature + +img_path = "D:/uni/master/machine_learning/StreetsignRecognition/data/processed" +vector_path = "D:/uni/master/machine_learning/StreetsignRecognition/data/vectors" + +def get_concept(path) -> Concept: + if "fahrtrichtung_links" in path: + return Concept.LINKS_ABBIEGEN + elif "fahrtrichtung_rechts" in path: + return Concept.RECHTS_ABBIEGEN + elif "rechts_vor_links" in path: + return Concept.RECHTS_VOR_LINKS + elif "stop" in path: + return Concept.STOP + elif "vorfahrt_gewaehren" in path: + return Concept.VORFAHRT_GEWAEHREN + elif "vorfahrtsstraße" in path: + return Concept.VORFAHRT_STRASSE + else: + return Concept.UNKNOWN + +if __name__ == "__main__": + # go to every file in img_path + for dirpath, dnames, fnames in os.walk(img_path): + save_path = os.path.join(dirpath.replace(img_path, vector_path)) + + # create directories if not exist + if not os.path.exists(save_path): + os.makedirs(save_path) + + for fname in fnames: + image_path = os.path.join(dirpath, fname) + try: + # extract features + image = cv.imread(image_path) + color_percentage_overall = get_color_percentage(image) + color_percentage_rasters = get_raster_color_percentage(image, 4, 4) + # add more + + except Exception: + print("Failed for ", image_path) + continue + + concept: Concept = get_concept(image_path) + if concept == Concept.UNKNOWN: + print("no concept found for ", image_path) + continue + + vec: FeatureVector = FeatureVector(concept) + vec.add_feature(Feature.OVERALL_COLOR_PERCENTAGE, color_percentage_overall) + vec.add_feature(Feature.RASTER_COLOR_PERCENTAGE, color_percentage_rasters) + # add more + print(vec.get_vector()) + vec_path = os.path.join(save_path, fname) + vec.save(vec_path) \ No newline at end of file