added color feature and raster_image | adjusted crop and remove_background

This commit is contained in:
Lukas Karras
2024-10-16 15:11:13 +02:00
parent c7ef220799
commit fd2a692281
4 changed files with 127 additions and 30 deletions

5
.gitignore vendored
View File

@@ -2,4 +2,7 @@
.vsc .vsc
.venv .venv
data data
# ignore all files in src that start with an underscore
src/**/_*

View File

@@ -1,8 +1,56 @@
class FeatureExtractor: import cv2 as cv
import numpy as np
def get_overall_color_percentage(): blue_lower_hsv = np.array([90, 75, 75])
pass blue_upper_hsv = np.array([130, 255, 255])
def get_color_per_raster_percentage(): red_lower1_hsv = np.array([0, 75, 75])
pass red_upper1_hsv = np.array([10, 255, 255])
red_lower2_hsv = np.array([170, 75, 75])
red_upper2_hsv = np.array([179, 255, 255])
yellow_lower_hsv = np.array([15, 75, 75])
yellow_upper_hsv = np.array([40, 255, 255])
white_lower_hsv = np.array([0, 0, 190])
white_upper_hsv = np.array([180, 65, 255])
black_lower_hsv = np.array([0, 0, 1]) # because 0 0 0 pixels are ignored as background
black_upper_hsv = np.array([180, 255, 30])
def get_color_percentage(image: cv.Mat) -> dict:
hsv = cv.cvtColor(image, cv.COLOR_BGR2HSV)
mask_red1 = cv.inRange(hsv, red_lower1_hsv, red_upper1_hsv)
mask_red2 = cv.inRange(hsv, red_lower2_hsv, red_upper2_hsv)
mask_red = mask_red1 | mask_red2
mask_blue = cv.inRange(hsv, blue_lower_hsv, blue_upper_hsv)
mask_yellow = cv.inRange(hsv, yellow_lower_hsv, yellow_upper_hsv)
mask_white = cv.inRange(hsv, white_lower_hsv, white_upper_hsv)
mask_black = cv.inRange(hsv, black_lower_hsv, black_upper_hsv)
width, height = image.shape[:2]
image_pixels = width*height
red_pixels = cv.countNonZero(mask_red)
blue_pixels = cv.countNonZero(mask_blue)
yellow_pixels = cv.countNonZero(mask_yellow)
white_pixels = cv.countNonZero(mask_white)
black_pixels = cv.countNonZero(mask_black)
red_percentage = int((red_pixels / image_pixels) * 100)
blue_percentage = int((blue_pixels / image_pixels) * 100)
yellow_percentage = int((yellow_pixels / image_pixels) * 100)
white_percentage = int((white_pixels / image_pixels) * 100)
black_percentage = int((black_pixels / image_pixels) * 100)
none_percentage = 100 - red_percentage - blue_percentage - yellow_percentage - white_percentage - black_percentage
return {
"red": red_percentage,
"blue": blue_percentage,
"yellow": yellow_percentage,
"white": white_percentage,
"black": black_percentage,
"none": none_percentage
}

View File

@@ -1,5 +1,8 @@
from typing import Any from typing import Any
import os
import pickle
from feature_vector import FeatureVector
from concept import Concept from concept import Concept
@@ -24,4 +27,18 @@ class FeatureVector:
def get_feature_value(self, key) -> Any: def get_feature_value(self, key) -> Any:
return self.features[key] return self.features[key]
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:
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

View File

@@ -1,14 +1,17 @@
import cv2 as cv import cv2 as cv
import numpy as np import numpy as np
def crop(image: cv.Mat) -> cv.Mat: def crop(image: cv.Mat, lower_rgb: np.ndarray, upper_rgb: np.ndarray) -> cv.Mat:
width, height = image.shape[:2] width, height = image.shape[:2]
x1 = None x1 = None
for x in range(width): for x in range(width):
for y in range(height): for y in range(height):
b, g, r = image[x][y] b, g, r = image[x][y]
if(r < 150 and g < 150 and b < 150): 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]
):
x1 = x x1 = x
break break
if x1 is not None: if x1 is not None:
@@ -18,7 +21,10 @@ def crop(image: cv.Mat) -> cv.Mat:
for y in range(height): for y in range(height):
for x in range(width): for x in range(width):
b, g, r = image[x][y] b, g, r = image[x][y]
if(r < 150 and g < 150 and b < 150): 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]
):
y1 = y y1 = y
break break
if y1 is not None: if y1 is not None:
@@ -28,7 +34,10 @@ def crop(image: cv.Mat) -> cv.Mat:
for x in range(width-1, 0, -1): for x in range(width-1, 0, -1):
for y in range(height): for y in range(height):
b, g, r = image[x][y] b, g, r = image[x][y]
if(r < 150 and g < 150 and b < 150): 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]
):
x2 = x x2 = x
break break
if x2 is not None: if x2 is not None:
@@ -38,7 +47,10 @@ def crop(image: cv.Mat) -> cv.Mat:
for y in range(height-1, 0, -1): for y in range(height-1, 0, -1):
for x in range(width): for x in range(width):
b, g, r = image[x][y] b, g, r = image[x][y]
if(r < 150 and g < 150 and b < 150): 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]
):
y2 = y y2 = y
break break
if y2 is not None: if y2 is not None:
@@ -47,26 +59,29 @@ def crop(image: cv.Mat) -> cv.Mat:
return image[x1-2:x2+2, y1-2:y2+2] return image[x1-2:x2+2, y1-2:y2+2]
def remove_background(image: cv.Mat, lower_rgb: np.ndarray) -> cv.Mat: def remove_background(image: cv.Mat, lower_rgb: np.ndarray, upper_rgb: np.ndarray) -> cv.Mat:
width, height = image.shape[:2] width, height = image.shape[:2]
# left to right # left to right
for x in range(width): for x in range(width):
for y in range(height): for y in range(height):
b, g, r = image[x][y] b, g, r = image[x][y]
if(r >= lower_rgb[0] and g >= lower_rgb[1] and b >= lower_rgb[2]): 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] = [0, 0, 0] image[x][y] = [0, 0, 0]
else: else:
break break
# top to bottom # # top to bottom
for y in range(height): # for y in range(height):
for x in range(width): # for x in range(width):
b, g, r = image[x][y] # b, g, r = image[x][y]
if(r >= lower_rgb[0] and g >= lower_rgb[1] and b >= lower_rgb[2]): # if(r >= lower_rgb[0] and g >= lower_rgb[1] and b >= lower_rgb[2]):
image[x][y] = [0, 0, 0] # image[x][y] = [0, 0, 0]
else: # else:
break # break
# right to left # right to left
for x in range(width): for x in range(width):
@@ -78,16 +93,30 @@ def remove_background(image: cv.Mat, lower_rgb: np.ndarray) -> cv.Mat:
break break
# bottom to top # bottom to top
for y in range(height): # for y in range(height):
for x in range(width-1, 0, -1): # for x in range(width-1, 0, -1):
b, g, r = image[x][y] # b, g, r = image[x][y]
if(r >= lower_rgb[0] and g >= lower_rgb[1] and b >= lower_rgb[2]): # if(r >= lower_rgb[0] and g >= lower_rgb[1] and b >= lower_rgb[2]):
image[x][y] = [0, 0, 0] # image[x][y] = [0, 0, 0]
else: # else:
break # break
return image return image
def raster_image(image: cv.Mat, num_cols: int = 4, num_rows: int = 4) -> dict[cv.Mat]: def raster_image(image: cv.Mat, num_cols: int = 4, num_rows: int = 4) -> list[list[cv.Mat]]:
pass width, height = image.shape[:2]
print("w", width, " h", height)
rasters = [[0 for x in range(num_rows)] for y in range(num_cols)]
print(rasters)
for y in range(num_cols):
for x in range(num_rows):
x1 = int(x * (width/num_rows))
y1 = int(y * (height/num_cols))
x2 = int((x+1) * (width/num_rows))
y2 = int((y+1) * (height/num_cols))
print(" x1", x1, " y1", y1, " x2", x2, " y2", y2)
rasters[x][y] = image[x1:x2, y1:y2]
return rasters