added program to create feature vectors

This commit is contained in:
Lukas K
2024-10-18 14:18:32 +02:00
parent c3d0e72067
commit ce2e7c863a
2 changed files with 62 additions and 0 deletions

View File

@@ -7,3 +7,4 @@ class Concept(Enum):
STOP = auto()
RECHTS_ABBIEGEN = auto()
LINKS_ABBIEGEN = auto()
RECHTS_VOR_LINKS = auto()

View File

@@ -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)