diff --git a/donkeycar/parts/cv.py b/donkeycar/parts/cv.py index 84f51e8be..31c5a0684 100644 --- a/donkeycar/parts/cv.py +++ b/donkeycar/parts/cv.py @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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: """ @@ -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 # @@ -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, @@ -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()