Skip to content

Commit

Permalink
fix cv color conversions (#1133)
Browse files Browse the repository at this point in the history 
- several of the color conversions were not correctly returning
  the converted image; return None instead.
- this commit fixes those and makes all color conversions
  similar in how they return that converted image
- This also adds a new transform; a version of the trapezoid
  mask that takes values from the edges of the image like
  the crop transform does (rather than absoluted image coordinates
  as the pre-existing trapezoid transform does).
  • Loading branch information
@Ezward
Ezward authored Jun 26, 2023
1 parent c0d4eb3 commit 3dee532
Showing 1 changed file with 78 additions and 22 deletions.
100 changes: 78 additions & 22 deletions donkeycar/parts/cv.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,7 @@ def run(self, img_arr):
return None

try:
img_arr = cv2.cvtColor(img_arr, cv2.COLOR_RGB2GRAY)
return img_arr
return cv2.cvtColor(img_arr, cv2.COLOR_RGB2GRAY)
except:
logger.error("Unable to convert RGB image to greyscale")
return None
Expand All @@ -40,7 +39,7 @@ def run(self, img_arr):
return None

try:
img_arr = cv2.cvtColor(img_arr, cv2.COLOR_GRAY2RGB)
return cv2.cvtColor(img_arr, cv2.COLOR_GRAY2RGB)
except:
logger.error(F"Unable to convert greyscale image of shape {img_arr.shape} to RGB")
return None
Expand All @@ -55,7 +54,7 @@ def run(self, img_arr):
return None

try:
img_arr = cv2.cvtColor(img_arr, cv2.COLOR_GRAY2BGR)
return cv2.cvtColor(img_arr, cv2.COLOR_GRAY2BGR)
except:
logger.error(F"Unable to convert greyscale image of shape {img_arr.shape} to RGB")
return None
Expand All @@ -75,8 +74,7 @@ def run(self, img_arr):
return None

try:
img_arr = cv2.cvtColor(img_arr, cv2.COLOR_BGR2GRAY)
return img_arr
return cv2.cvtColor(img_arr, cv2.COLOR_BGR2GRAY)
except:
logger.error("Unable to convert BGR image to greyscale")
return None
Expand All @@ -92,8 +90,7 @@ def run(self, img_arr):
return None

try:
img_arr = cv2.cvtColor(img_arr, cv2.COLOR_HSV2GRAY)
return img_arr
return cv2.cvtColor(img_arr, cv2.COLOR_HSV2GRAY)
except:
logger.error("Unable to convert HSV image to greyscale")
return None
Expand Down Expand Up @@ -121,8 +118,7 @@ def run(self, img_arr):
return None

try:
img_arr = cv2.cvtColor(img_arr, cv2.COLOR_BGR2RGB)
return img_arr
return cv2.cvtColor(img_arr, cv2.COLOR_BGR2RGB)
except:
logger.error("Unable to convert BGR image to RGB")
return None
Expand All @@ -138,8 +134,7 @@ def run(self, img_arr):
return None

try:
img_arr = cv2.cvtColor(img_arr, cv2.COLOR_RGB2BGR)
return img_arr
return cv2.cvtColor(img_arr, cv2.COLOR_RGB2BGR)
except:
logger.error("Unable to convert RGB image to BRG")
return None
Expand All @@ -155,8 +150,7 @@ def run(self, img_arr):
return None

try:
img_arr = cv2.cvtColor(img_arr, cv2.COLOR_HSV2RGB)
return img_arr
return cv2.cvtColor(img_arr, cv2.COLOR_HSV2RGB)
except:
logger.error("Unable to convert HSV image to RGB")
return None
Expand All @@ -172,8 +166,7 @@ def run(self, img_arr):
return None

try:
img_arr = cv2.cvtColor(img_arr, cv2.COLOR_RGB2HSV)
return img_arr
return cv2.cvtColor(img_arr, cv2.COLOR_RGB2HSV)
except:
logger.error("Unable to convert RGB image to HSV")
return None
Expand All @@ -189,8 +182,7 @@ def run(self, img_arr):
return None

try:
img_arr = cv2.cvtColor(img_arr, cv2.COLOR_HSV2BGR)
return img_arr
return cv2.cvtColor(img_arr, cv2.COLOR_HSV2BGR)
except:
logger.error("Unable to convert HSV image to BGR")
return None
Expand All @@ -206,8 +198,7 @@ def run(self, img_arr):
return None

try:
img_arr = cv2.cvtColor(img_arr, cv2.COLOR_BGR2HSV)
return img_arr
return cv2.cvtColor(img_arr, cv2.COLOR_BGR2HSV)
except:
logger.error("Unable to convert BGR image to HSV")
return None
Expand Down Expand Up @@ -422,6 +413,62 @@ def shutdown(self):
self.masks = {} # free cached masks


class ImgTrapezoidalEdgeMask:
def __init__(self, upper_left, upper_right, lower_left, lower_right, top, bottom, fill=[255,255,255]) -> None:
"""
Apply a trapezoidal mask to an image, where bounds are
relative the the edge of the image, conserving the
image pixels within the trapezoid and masking everything
other pixels with the fill color
"""
self.lower_left = lower_left
self.lower_right = lower_right
self.upper_left = upper_left
self.upper_right = upper_right
self.top = top
self.bottom = bottom
self.fill = fill
self.masks = {}

def run(self, image):
"""
Apply trapezoidal mask
# # # # # # # # # # # # #
# xxxxxxxxxxxxxxxxxxxxxxx
# xxxx ul ur xxxxxxxx min_y
# xxx xxxxxxx
# xx xxxxxx
# x xxxxx
# ll lr xx max_y
"""
transformed = None
if image is not None:
mask = None
key = str(image.shape)
if self.masks.get(key) is None:
height, width, depth = image_shape(image)
mask = np.zeros(image.shape, dtype=np.int32)
points = [
[self.upper_left, self.top],
[width - self.upper_right, self.top],
[width - self.lower_right, height - self.bottom],
[self.lower_left, height - self.bottom]
]
cv2.fillConvexPoly(mask,
np.array(points, dtype=np.int32),
self.fill)
mask = np.asarray(mask, dtype='bool')
self.masks[key] = mask

mask = self.masks[key]
transformed = np.multiply(image, mask)

return transformed

def shutdown(self):
self.masks = {} # free cached masks


class ImgCropMask:
def __init__(self, left=0, top=0, right=0, bottom=0, fill=[255, 255, 255]) -> None:
"""
Expand Down Expand Up @@ -746,7 +793,7 @@ def shutdown(self):
#
# masking tranformations
#
if "TRAPEZE" == transformation or "CROP" == transformation:
if "TRAPEZE" == transformation or "TRAPEZE_EDGE" == transformation or "CROP" == transformation:
#
# masking transformations
#
Expand All @@ -759,6 +806,15 @@ def shutdown(self):
args.top if args.top is not None else 0,
args.bottom if args.bottom is not None else height
)
elif "TRAPEZE_EDGE" == transformation:
transformer = ImgTrapezoidalEdgeMask(
args.left if args.left is not None else 0,
args.right if args.right is not None else width,
args.left_bottom if args.left_bottom is not None else 0,
args.right_bottom if args.right_bottom is not None else width,
args.top if args.top is not None else 0,
args.bottom if args.bottom is not None else height
)
else:
transformer = ImgCropMask(
args.left if args.left is not None else 0,
Expand All @@ -782,7 +838,7 @@ def shutdown(self):
transformer = ImgHSV2BGR()
elif "RGB2GREY" == transformation:
transformer = ImgRGB2GRAY()
elif "RBGR2GREY" == transformation:
elif "BGR2GREY" == transformation:
transformer = ImgBGR2GRAY()
elif "HSV2GREY" == transformation:
transformer = ImgHSV2GRAY()
Expand Down

0 comments on commit 3dee532

Please sign in to comment.