From 2484e8f8540708103dd9f038e666cbd59a33ebe1 Mon Sep 17 00:00:00 2001 From: Lukas K Date: Fri, 18 Oct 2024 11:25:35 +0200 Subject: [PATCH] added feature enum | fixed image_processing crop --- src/feature.py | 7 ++++++ src/feature_vector.py | 20 +++++++++-------- src/image_processing.py | 49 ++++++++++++++++++++++++----------------- 3 files changed, 47 insertions(+), 29 deletions(-) create mode 100644 src/feature.py diff --git a/src/feature.py b/src/feature.py new file mode 100644 index 0000000..27a34bf --- /dev/null +++ b/src/feature.py @@ -0,0 +1,7 @@ +from enum import Enum, auto + +class Feature(Enum): + OVERALL_COLOR_PERCENTAGE = auto() + RASTER_COLOR_PERCENTAGE = auto() + CORNERS = auto() + EDGES = auto() diff --git a/src/feature_vector.py b/src/feature_vector.py index 2aa7f0f..0077e9c 100644 --- a/src/feature_vector.py +++ b/src/feature_vector.py @@ -1,9 +1,9 @@ -from typing import Any +from typing import Any, Self import os import pickle -from feature_vector import FeatureVector from concept import Concept +from feature import Feature class FeatureVector: @@ -13,7 +13,7 @@ class FeatureVector: self.concept: Concept = concept - def add_feature(self, key, value) -> None: + def add_feature(self, key: Feature, value: Any) -> None: self.features.update({key: value}) @@ -25,20 +25,22 @@ class FeatureVector: return len(self.features) - def get_feature_value(self, key) -> Any: + def get_feature_value(self, key: str) -> Any: return self.features[key] + def get_features(self) -> Any: + return self.features + + def save(self, path: os.path) -> None: with open(path, "wb") as write: pickle.dump(self, write, pickle.HIGHEST_PROTOCOL) - def load(self, path: os.path) -> FeatureVector: + def load(path: os.path) -> Self: with open(path, "rb") as read: - feature_vector: FeatureVector = pickle.load(read) - self.features = feature_vector.features - self.concept = feature_vector.concept - return feature_vector + self: FeatureVector = pickle.load(read) + return self diff --git a/src/image_processing.py b/src/image_processing.py index e6e4304..e485e5c 100644 --- a/src/image_processing.py +++ b/src/image_processing.py @@ -8,9 +8,9 @@ def crop(image: cv.Mat, lower_rgb: np.ndarray, upper_rgb: np.ndarray) -> cv.Mat: for x in range(width): for y in range(height): b, g, r = image[x][y] - if(r <= lower_rgb[0] and r >= upper_rgb[0] and - g <= lower_rgb[1] and g >= upper_rgb[1] and - b <= lower_rgb[2] and b >= upper_rgb[2] + if(not (r >= lower_rgb[0] and r <= upper_rgb[0] and + g >= lower_rgb[1] and g <= upper_rgb[1] and + b >= lower_rgb[2] and b <= upper_rgb[2]) ): x1 = x break @@ -21,9 +21,9 @@ def crop(image: cv.Mat, lower_rgb: np.ndarray, upper_rgb: np.ndarray) -> cv.Mat: for y in range(height): for x in range(width): b, g, r = image[x][y] - if(r <= lower_rgb[0] and r >= upper_rgb[0] and - g <= lower_rgb[1] and g >= upper_rgb[1] and - b <= lower_rgb[2] and b >= upper_rgb[2] + if(not (r >= lower_rgb[0] and r <= upper_rgb[0] and + g >= lower_rgb[1] and g <= upper_rgb[1] and + b >= lower_rgb[2] and b <= upper_rgb[2]) ): y1 = y break @@ -34,9 +34,9 @@ def crop(image: cv.Mat, lower_rgb: np.ndarray, upper_rgb: np.ndarray) -> cv.Mat: for x in range(width-1, 0, -1): for y in range(height): b, g, r = image[x][y] - if(r <= lower_rgb[0] and r >= upper_rgb[0] and - g <= lower_rgb[1] and g >= upper_rgb[1] and - b <= lower_rgb[2] and b >= upper_rgb[2] + if(not (r >= lower_rgb[0] and r <= upper_rgb[0] and + g >= lower_rgb[1] and g <= upper_rgb[1] and + b >= lower_rgb[2] and b <= upper_rgb[2]) ): x2 = x break @@ -47,9 +47,9 @@ def crop(image: cv.Mat, lower_rgb: np.ndarray, upper_rgb: np.ndarray) -> cv.Mat: for y in range(height-1, 0, -1): for x in range(width): b, g, r = image[x][y] - if(r <= lower_rgb[0] and r >= upper_rgb[0] and - g <= lower_rgb[1] and g >= upper_rgb[1] and - b <= lower_rgb[2] and b >= upper_rgb[2] + if(not (r >= lower_rgb[0] and r >= upper_rgb[0] and + g >= lower_rgb[1] and g >= upper_rgb[1] and + b >= lower_rgb[2] and b >= upper_rgb[2]) ): y2 = y break @@ -59,7 +59,7 @@ def crop(image: cv.Mat, lower_rgb: np.ndarray, upper_rgb: np.ndarray) -> cv.Mat: return image[x1-2:x2+2, y1-2:y2+2] -def remove_background(image: cv.Mat, lower_rgb: np.ndarray, upper_rgb: np.ndarray) -> cv.Mat: +def remove_background(image: cv.Mat, lower_rgb: np.ndarray, upper_rgb: np.ndarray, new_background: np.ndarray) -> cv.Mat: width, height = image.shape[:2] # left to right @@ -70,7 +70,7 @@ def remove_background(image: cv.Mat, lower_rgb: np.ndarray, upper_rgb: np.ndarra g >= lower_rgb[1] and g <= upper_rgb[1] and b >= lower_rgb[2] and b <= upper_rgb[2] ): - image[x][y] = [0, 0, 0] + image[x][y] = new_background else: break @@ -78,8 +78,11 @@ def remove_background(image: cv.Mat, lower_rgb: np.ndarray, upper_rgb: np.ndarra # for y in range(height): # for x in range(width): # b, g, r = image[x][y] - # if(r >= lower_rgb[0] and g >= lower_rgb[1] and b >= lower_rgb[2]): - # image[x][y] = [0, 0, 0] + # if(r >= lower_rgb[0] and r <= upper_rgb[0] and + # g >= lower_rgb[1] and g <= upper_rgb[1] and + # b >= lower_rgb[2] and b <= upper_rgb[2] + # ): + # image[x][y] = new_background # else: # break @@ -87,8 +90,11 @@ def remove_background(image: cv.Mat, lower_rgb: np.ndarray, upper_rgb: np.ndarra for x in range(width): for y in range(height-1, 0, -1): b, g, r = image[x][y] - if(r >= lower_rgb[0] and g >= lower_rgb[1] and b >= lower_rgb[2]): - image[x][y] = [0, 0, 0] + if(r >= lower_rgb[0] and r <= upper_rgb[0] and + g >= lower_rgb[1] and g <= upper_rgb[1] and + b >= lower_rgb[2] and b <= upper_rgb[2] + ): + image[x][y] = new_background else: break @@ -96,8 +102,11 @@ def remove_background(image: cv.Mat, lower_rgb: np.ndarray, upper_rgb: np.ndarra # for y in range(height): # for x in range(width-1, 0, -1): # b, g, r = image[x][y] - # if(r >= lower_rgb[0] and g >= lower_rgb[1] and b >= lower_rgb[2]): - # image[x][y] = [0, 0, 0] + # if(r >= lower_rgb[0] and r <= upper_rgb[0] and + # g >= lower_rgb[1] and g <= upper_rgb[1] and + # b >= lower_rgb[2] and b <= upper_rgb[2] + # ): + # image[x][y] = new_background # else: # break