def _find_batteries(im): blue = extract_color(im, 216 / 2, (100, 255), (50, 150)) green = extract_color(im, 105 / 2, (200, 255), (200, 255)) color = blue + green structuring_element1 = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15)) # structuring_element2 = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5)) color = cv2.morphologyEx(color, cv2.MORPH_CLOSE, structuring_element1) # color = cv2.morphologyEx(color, cv2.MORPH_OPEN, structuring_element2) # color = scale(color, 0.25) # im = scale(im, 0.25) # show(color) contours, _ = cv2.findContours(color.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # show(get_drawn_contours(color, contours, True)) contours = [contour_bounding_box_for_contour(c) for c in contours] height, width = im.shape[:2] def is_large_enough(contour): x, y, w, h = cv2.boundingRect(contour) return w >= width * 0.1 and h >= height * 0.1 contours = [c for c in contours if is_large_enough(c)] return [four_point_transform(im, c) for c in contours]
def get_connections(im): """ Returns a list of tuples, each with (color, destination, (to_click_x, to_click_y)). The list is ordered from top to bottom. """ blue = extract_color(im, 115, (100, 255), (0, 255)) red = extract_color(im, 5, (100, 255), (0, 255)) black = extract_color(im, (0, 255), (0, 255), (0, 5)) colors = ( (Color.BLUE, blue), (Color.BLACK, black), (Color.RED, red), ) mask = blue.copy() height, width = mask.shape rows = (0.33, 0.48, 0.62) cols = (0.30, 0.55) row_pxs = [int(row_percent * height) for row_percent in rows] start_col_px, end_col_px = [int(col_percent * width) for col_percent in cols] def _get_output_for_connection(color, start_row_idx, end_row_idx): """Returns a tuple with (color, destination, (to_click_x, to_click_y))""" # Use the start_row to find a point to click to cut the wire. to_click = (start_col_px, row_pxs[start_row_idx]) # Figure out which letter we're connecting to based on the end row. destination = ORDERED_DESTINATIONS[end_row_idx] return color, destination, to_click output = [] for start_row, start_row_px in enumerate(row_pxs): current_output = None for end_row, end_row_px in enumerate(row_pxs): # Reset the mask mask[:] = 0 # Draw a line where the connection should be on the mask cv2.line(mask, (start_col_px, start_row_px), (end_col_px, end_row_px), 255, 15) # Convert the mask to be usable in a masked_array bool_mask = np.invert(mask.astype(bool)) # We only want to consider it a wire there if there's at least 75% of the mask filled in activation_threshold = bool_mask.sum() * .66 for color, color_mat in colors: activated_amount = np.ma.masked_array(color_mat, mask=bool_mask).sum() if activated_amount > activation_threshold: new_output = _get_output_for_connection(color, start_row, end_row) new_output_and_activated = (new_output, activated_amount) if current_output is None: current_output = new_output_and_activated else: current_output = sorted((current_output, new_output_and_activated), key=lambda x: x[1])[-1] # combined = np.zeros_like(im) # combined[:, :, 0] = color_mat # combined[:, :, 2] = mask # show(combined) if current_output is not None: print current_output output.append(current_output[0]) return output
def _get_button_images(im): im_mono = extract_color(im, (0, 180), (0, 120), (0, 120)) contours, hierarchy = cv2.findContours(im_mono, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) contours = [cv2.approxPolyDP(c, 6, True) for c in contours] midpoint_y = im.shape[1] / 2 contours = [c for c in contours if len(c) == 4 and cv2.boundingRect(c)[1] > midpoint_y and cv2.isContourConvex(c)] button_box = sorted(contours, key=cv2.contourArea)[-1] button_box = button_box.reshape((4, 2)) button_im = four_point_transform(im, button_box) button_im_mono = extract_color(button_im, 18, (50, 100), (175, 255)) contours, hierarchy = cv2.findContours(button_im_mono, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # There's 4 buttons, which will be the largest 4 contours contours = sorted(contours, key=cv2.contourArea)[-4:] contour_rects = [cv2.boundingRect(c) for c in contours] # Sort from left to right contour_rects = sorted(contour_rects, key=lambda rect: rect[0]) buttons = [] for x, y, w, h in contour_rects: button = four_point_transform(button_im, np.array(((x, y), (x + w, y), (x, y + h), (x + w, y + h))), -6) button = extract_color(button, 18, (50, 100), (175, 255)) buttons.append(button) return buttons
def _get_button_images(im): im_mono = extract_color(im, (0, 180), (0, 120), (0, 120)) contours, hierarchy = cv2.findContours(im_mono, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) contours = [cv2.approxPolyDP(c, 6, True) for c in contours] midpoint_y = im.shape[1] / 2 contours = [ c for c in contours if len(c) == 4 and cv2.boundingRect(c)[1] > midpoint_y and cv2.isContourConvex(c) ] button_box = sorted(contours, key=cv2.contourArea)[-1] button_box = button_box.reshape((4, 2)) button_im = four_point_transform(im, button_box) button_im_mono = extract_color(button_im, 18, (50, 100), (175, 255)) contours, hierarchy = cv2.findContours(button_im_mono, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # There's 4 buttons, which will be the largest 4 contours contours = sorted(contours, key=cv2.contourArea)[-4:] contour_rects = [cv2.boundingRect(c) for c in contours] # Sort from left to right contour_rects = sorted(contour_rects, key=lambda rect: rect[0]) buttons = [] for x, y, w, h in contour_rects: button = four_point_transform( button_im, np.array(((x, y), (x + w, y), (x, y + h), (x + w, y + h))), -6) button = extract_color(button, 18, (50, 100), (175, 255)) buttons.append(button) return buttons
def _get_button_text(im, tesseract): im = get_subset(im, _TEXT_X_PERCENTS, _TEXT_Y_PERCENTS) black_text = extract_color(im, (0, 255), (0, 255), (0, 65)) white_text = extract_color(im, (0, 255), (0, 40), (230, 255)) if black_text.any(): text_image = black_text else: assert white_text.any(), "Neither black nor white text have any pixels." text_image = white_text tesseract.SetImage(Image.fromarray(text_image)) word = tesseract.GetUTF8Text().strip() # It messes up a lot, just take the one with the closest edit distance return sorted(ButtonLabel, key=lambda s: editdistance.eval(s.value, word))[0]
def _get_button_text(im, tesseract): im = get_subset(im, _TEXT_X_PERCENTS, _TEXT_Y_PERCENTS) black_text = extract_color(im, (0, 255), (0, 255), (0, 65)) white_text = extract_color(im, (0, 255), (0, 40), (230, 255)) if black_text.any(): text_image = black_text else: assert white_text.any( ), "Neither black nor white text have any pixels." text_image = white_text tesseract.SetImage(Image.fromarray(text_image)) word = tesseract.GetUTF8Text().strip() # It messes up a lot, just take the one with the closest edit distance return sorted(ButtonLabel, key=lambda s: editdistance.eval(s.value, word))[0]
def _get_screen_image(im): im_mono = extract_color(im, 76, (50, 150), (50, 150)) # show(im_mono) contours, hierarchy = cv2.findContours(im_mono, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) midpoint_y = im.shape[1] * 2 / 3 contours_filtered = [] for contour in contours: # contour = cv2.approxPolyDP(contour, 2, True) x, y, w, h = cv2.boundingRect(contour) # if y + h < midpoint_y and cv2.isContourConvex(contour): if y + h < midpoint_y: contours_filtered.append(contour) contours = sorted(contours_filtered, key=cv2.contourArea)[-4:] boxes = [cv2.boundingRect(c) for c in contours] x1 = min(x for x, y, w, h in boxes) y1 = min(y for x, y, w, h in boxes) x2 = max(x + w for x, y, w, h in boxes) y2 = max(y + h for x, y, w, h in boxes) points = np.array(((x1, y1), (x1, y2), (x2, y1), (x2, y2))) screen = four_point_transform(im, points) screen_mono = cv2.cvtColor(screen, cv2.COLOR_BGR2GRAY) _, screen_mono = cv2.threshold(screen_mono, 245, 255, cv2.THRESH_BINARY) return screen_mono
def get_button_locations(im): mono = extract_color(im, 120, (0, 100), (0, 100)) contours, hierarchy = cv2.findContours(mono, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) contours = [cv2.approxPolyDP(contour, 6, True) for contour in contours] height, width = im.shape[:2] def is_in_middle_quarter_vertically_or_horizontally(contour_to_check): x, y, w, h = cv2.boundingRect(contour_to_check) return (y > height / 4 and y + h < height * 3 / 4) or \ (x > width / 4 and x + w < width * 3 / 4) contours = [c for c in contours if is_in_middle_quarter_vertically_or_horizontally(c)] contours = sorted(contours, key=cv2.contourArea, reverse=True)[:4] assert len(contours) == 4, "Expected to find 4 buttons, found %s" % len(contours) centers = [] for c in contours: x, y, w, h = cv2.boundingRect(c) centers.append((x + (w / 2), y + (h / 2))) x_sort = sorted(centers, key=lambda center: center[0]) y_sort = sorted(centers, key=lambda center: center[1]) top = y_sort[0] bottom = y_sort[-1] left = x_sort[0] right = x_sort[-1] assert top != bottom != left != right, "Expected each point to be different" return top, right, bottom, left
def _get_leds_are_lit(im): leds_are_lit = [] for i in range(len(_TOP_X_BOUNDARIES) - 1): led = get_subset(im, _TOP_X_BOUNDARIES[i:i + 2], _LED_Y_BOUNDARIES) lit_led = extract_color(led, 51 / 2, (40, 90), (220, 255)) leds_are_lit.append(lit_led.any()) # show(lit_led) return leds_are_lit
def _get_wire_positions(im, wire_color, hue, saturation=(150, 255), value=(100, 255)): color = extract_color(im, hue, saturation, value) contours = get_contours(color, close_and_open=False) return [(get_center_for_contour(c), wire_color) for c in contours]
def _get_button_locations(im): im_mono = extract_color(im, 18, (50, 100), (175, 255)) contours, hierarchy = cv2.findContours(im_mono, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # There's two buttons, which will be the largest two contours contours = sorted(contours, key=cv2.contourArea)[-4:] centers = [get_center_for_contour(c) for c in contours] # Sort them from left to right centers = sorted(centers, key=lambda center: center[0]) return centers
def get_down_button(im): im_mono = extract_color(im, 18, (50, 100), (175, 255)) contours, hierarchy = cv2.findContours(im_mono, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # There's two buttons, which will be the largest two contours contours = sorted(contours, key=cv2.contourArea)[-2:] centers = [get_center_for_contour(c) for c in contours] # We want the lower one on screen, which is the larger y value centers = sorted(centers, key=lambda center: center[1]) return centers[-1]
def get_strip_color(im): im = get_subset(im, _STRIP_X_PERCENTS, _STRIP_Y_PERCENTS) colors = [color for mat, color in ( (extract_color(im, (0, 180), (0, 5), (190, 255)), StripColor.WHITE), (extract_color_2(im, 50, 94, 84), StripColor.YELLOW), (extract_color_2(im, 0, 82, 76), StripColor.RED), (extract_color_2(im, 218, 85, 79), StripColor.BLUE), ) if mat.any()] assert len(colors) == 1, "Strip does not look like one color" return colors[0]
def _extract_side(im, is_bottom): if is_bottom: color = extract_color(im, 32 / 2, (120, 255), (100, 220)) else: color = extract_color(im, 32 / 2, (100, 255), (100, 255)) structuring_element1 = cv2.getStructuringElement(cv2.MORPH_RECT, (10, 10)) structuring_element2 = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5)) color = cv2.morphologyEx(color, cv2.MORPH_CLOSE, structuring_element1) color = cv2.morphologyEx(color, cv2.MORPH_OPEN, structuring_element2) # show(im) # show(color) height, width = color.shape[:2] contours, _ = cv2.findContours(color.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # contours = [cv2.approxPolyDP(c, 0.03 * cv2.arcLength(c, True), True) for c in contours] contours = [ cv2.approxPolyDP(c, 0.015 * cv2.arcLength(c, True), True) for c in contours ] # show(get_drawn_contours(c2, contours, True)) # for c in contours: # print len(c) # show(get_drawn_contours(c2, [c], True)) contours = [c for c in contours if len(c) == 4] contours = sorted(contours, key=cv2.contourArea, reverse=True)[:2] a, b = contours points = list(a) + list(b) tl = point_closest_to(points, 0, 0) tr = point_closest_to(points, width, 0) bl = point_closest_to(points, 0, height) br = point_closest_to(points, width, height) contour = np.array([tl, tr, br, bl]) contour = contour.reshape((4, 2)) # show(get_drawn_contours(c2, [contour], True)) return four_point_transform(im, contour, margin_percent=5)
def _get_has_stars(im): has_stars = [] for i in range(len(_BOTTOM_X_BOUNDARIES) - 1): star = get_subset(im, _BOTTOM_X_BOUNDARIES[i:i + 2], _STAR_Y_BOUNDARIES) has_star = extract_color(star, 33 / 2, (75, 125), (0, 70)) # show(has_star) w, h = get_dimens(star) star_ratio = float(cv2.countNonZero(has_star)) / (w * h) # print star_ratio has_stars.append(star_ratio > _STAR_RATIO_THRESHOLD) return has_stars
def get_clock_time_from_full_screenshot(full_screenshot, current_module_position, screenshot_helper): clock_im = _get_clock_image_from_full_screenshot(full_screenshot, current_module_position, screenshot_helper) assert clock_im is not None, "Unable to find clock" clock_bg = extract_color(clock_im, (0, 180), (0, 255), (0, 50)) contours = [ c for c in simplify_contours(get_contours(clock_bg), 0.001) if len(c) == 4 ] contour = max(contours, key=cv2.contourArea) display = four_point_transform(clock_im, contour) digits_im = extract_color(display, 0, (250, 255), (250, 255)) digit_images = [ get_subset(digits_im, _CLOCK_DIGIT_X_PERCENTS[digit_index:digit_index + 2], (0, 100)) for digit_index in range(len(_CLOCK_DIGIT_X_PERCENTS) - 1) ] # Determine if there is a colon or period separator. is_colon = get_subset(digit_images[2], (0, 100), (0, 50)).any() if is_colon: minutes_images = digit_images[:2] seconds_images = digit_images[-2:] else: minutes_images = [] seconds_images = digit_images[:2] minutes_digits = [_get_digit_from_image(im) for im in minutes_images] seconds_digits = [_get_digit_from_image(im) for im in seconds_images] if not minutes_digits: minutes_digits = [0] return minutes_digits, seconds_digits
def get_strip_color(im): im = get_subset(im, _STRIP_X_PERCENTS, _STRIP_Y_PERCENTS) colors = [ color for mat, color in ( (extract_color(im, (0, 180), (0, 5), (190, 255)), StripColor.WHITE), (extract_color_2(im, 50, 94, 84), StripColor.YELLOW), (extract_color_2(im, 0, 82, 76), StripColor.RED), (extract_color_2(im, 218, 85, 79), StripColor.BLUE), ) if mat.any() ] assert len(colors) == 1, "Strip does not look like one color" return colors[0]
def _extract_side(im, is_bottom): if is_bottom: color = extract_color(im, 32 / 2, (120, 255), (100, 220)) else: color = extract_color(im, 32 / 2, (100, 255), (100, 255)) structuring_element1 = cv2.getStructuringElement(cv2.MORPH_RECT, (10, 10)) structuring_element2 = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5)) color = cv2.morphologyEx(color, cv2.MORPH_CLOSE, structuring_element1) color = cv2.morphologyEx(color, cv2.MORPH_OPEN, structuring_element2) # show(im) # show(color) height, width = color.shape[:2] contours, _ = cv2.findContours(color.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # contours = [cv2.approxPolyDP(c, 0.03 * cv2.arcLength(c, True), True) for c in contours] contours = [cv2.approxPolyDP(c, 0.015 * cv2.arcLength(c, True), True) for c in contours] # show(get_drawn_contours(c2, contours, True)) # for c in contours: # print len(c) # show(get_drawn_contours(c2, [c], True)) contours = [c for c in contours if len(c) == 4] contours = sorted(contours, key=cv2.contourArea, reverse=True)[:2] a, b = contours points = list(a) + list(b) tl = point_closest_to(points, 0, 0) tr = point_closest_to(points, width, 0) bl = point_closest_to(points, 0, height) br = point_closest_to(points, width, height) contour = np.array([tl, tr, br, bl]) contour = contour.reshape((4, 2)) # show(get_drawn_contours(c2, [contour], True)) return four_point_transform(im, contour, margin_percent=5)
def get_has_parallel_port_for_side(side): side_w, side_h = get_dimens(side) if side_h > side_w: side = cv2.transpose(side) side_w, side_h = get_dimens(side) color = extract_color(side, 337 / 2, (75, 175), (175, 255)) contours = sorted(get_contours(color), key=cv2.contourArea) if not contours: return False contour = contours[-1] x, y, w, h = cv2.boundingRect(contour) width_percent = float(w) / side_w height_percent = float(h) / side_h return width_percent > _WIDTH_THRESHOLD and height_percent > _HEIGHT_THRESHOLD
def _get_count_for_subsection(battery): b_height, b_width = battery.shape[:2] if b_height > b_width: b_height, b_width = b_width, b_height # Rotate battery to be left-right oriented battery = cv2.transpose(battery) blue = extract_color(battery, 216 / 2, (100, 255), (50, 150)) cropped = blue[b_height/3:2*b_height/3, b_width/3:2*b_width/3] if cropped.any(): return 2 else: return 1
def get_clock_time_from_full_screenshot(full_screenshot, current_module_position, screenshot_helper): clock_im = _get_clock_image_from_full_screenshot(full_screenshot, current_module_position, screenshot_helper) assert clock_im is not None, "Unable to find clock" clock_bg = extract_color(clock_im, (0, 180), (0, 255), (0, 50)) contours = [c for c in simplify_contours(get_contours(clock_bg), 0.001) if len(c) == 4] contour = max(contours, key=cv2.contourArea) display = four_point_transform(clock_im, contour) digits_im = extract_color(display, 0, (250, 255), (250, 255)) digit_images = [ get_subset(digits_im, _CLOCK_DIGIT_X_PERCENTS[digit_index:digit_index + 2], (0, 100)) for digit_index in range(len(_CLOCK_DIGIT_X_PERCENTS) - 1) ] # Determine if there is a colon or period separator. is_colon = get_subset(digit_images[2], (0, 100), (0, 50)).any() if is_colon: minutes_images = digit_images[:2] seconds_images = digit_images[-2:] else: minutes_images = [] seconds_images = digit_images[:2] minutes_digits = [_get_digit_from_image(im) for im in minutes_images] seconds_digits = [_get_digit_from_image(im) for im in seconds_images] if not minutes_digits: minutes_digits = [0] return minutes_digits, seconds_digits
def _get_largest_contour_area_ratio_for_color(im, color): hue_out_of_360 = { LitSquare.YELLOW: 60, LitSquare.BLUE: 193, LitSquare.RED: 19, LitSquare.GREEN: 149, }[color] color_mat = extract_color(im, hue_out_of_360 / 2, (100, 255), (225, 255)) contour_areas = [cv2.contourArea(c) for c in get_contours(color_mat)] if contour_areas: largest_area = sorted(contour_areas)[-1] else: largest_area = 0 w, h = get_dimens(im) im_area = w * h area_ratio = float(largest_area) / im_area return area_ratio, color
def _get_largest_contour_area_ratio_for_color(im, color): hue_out_of_360 = { LitSquare.YELLOW: 60, LitSquare.BLUE: 193, LitSquare.RED: 19, LitSquare.GREEN: 149, }[color] color_mat = extract_color(im, hue_out_of_360/2, (100, 255), (225, 255)) contour_areas = [cv2.contourArea(c) for c in get_contours(color_mat)] if contour_areas: largest_area = sorted(contour_areas)[-1] else: largest_area = 0 w, h = get_dimens(im) im_area = w * h area_ratio = float(largest_area) / im_area return area_ratio, color
def find_arrows(im): mono = extract_color(im, 120, (50, 100), (50, 100)) contours, hierarchy = cv2.findContours(mono, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) contours = [cv2.approxPolyDP(contour, 6, True) for contour in contours] one_quarter_height = im.shape[0] / 4 three_quarter_height = im.shape[0] * 3 / 4 def is_in_middle_quarter_vertically(contour_to_check): x, y, w, h = cv2.boundingRect(contour_to_check) return y > one_quarter_height and y + h < three_quarter_height contours = [c for c in contours if cv2.isContourConvex(c) and is_in_middle_quarter_vertically(c)] contours = sorted(contours, key=cv2.contourArea, reverse=True)[:2] centers = [get_center_for_contour(c) for c in contours] centers = sorted(centers, key=lambda center: center[0]) return centers
def get_button_locations(im): mono = extract_color(im, 120, (0, 100), (0, 100)) contours, hierarchy = cv2.findContours(mono, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) contours = [cv2.approxPolyDP(contour, 6, True) for contour in contours] height, width = im.shape[:2] def is_in_middle_quarter_vertically_or_horizontally(contour_to_check): x, y, w, h = cv2.boundingRect(contour_to_check) return (y > height / 4 and y + h < height * 3 / 4) or \ (x > width / 4 and x + w < width * 3 / 4) contours = [ c for c in contours if is_in_middle_quarter_vertically_or_horizontally(c) ] contours = sorted(contours, key=cv2.contourArea, reverse=True)[:4] assert len( contours) == 4, "Expected to find 4 buttons, found %s" % len(contours) centers = [] for c in contours: x, y, w, h = cv2.boundingRect(c) centers.append((x + (w / 2), y + (h / 2))) x_sort = sorted(centers, key=lambda center: center[0]) y_sort = sorted(centers, key=lambda center: center[1]) top = y_sort[0] bottom = y_sort[-1] left = x_sort[0] right = x_sort[-1] assert top != bottom != left != right, "Expected each point to be different" return top, right, bottom, left
def _get_indicator_images_and_light_statuses(im): red = extract_color(im, 0, (50, 200), (50, 200)) # show(red) w_total, h_total = get_dimens(im) w_threshold = int(_INDICATOR_WIDTH_PERCENT_THRESHOLD * w_total) h_threshold = int(_INDICATOR_HEIGHT_PERCENT_THRESHOLD * h_total) def is_indicator_big_enough(contour): _, _, contour_w, contour_h = cv2.boundingRect(contour) return contour_w > w_threshold and contour_h > h_threshold contours = [ contour_bounding_box_for_contour(c) for c in get_contours(red) if is_indicator_big_enough(c) ] indicators = [four_point_transform(im, c) for c in contours] indicators_and_lights = [] for indicator in indicators: w, h = get_dimens(indicator) if w < h: # Rotate 90 degrees so it's horizontal indicator = rotate_image_clockwise(indicator) w, h = get_dimens(indicator) # Check if light is on left or right, flip accordingly light_width_threshold = w * _LIGHT_WIDTH_THRESHOLD light_height_threshold = h * _LIGHT_HEIGHT_THRESHOLD light_on = extract_color(indicator, (0, 180), (0, 0), (255, 255)) light_off = extract_color(indicator, (0, 180), (0, 40), (0, 50)) # show(light_on) # show(light_off) def is_light_big_enough(contour): _, _, contour_w, contour_h = cv2.boundingRect(contour) return contour_w > light_width_threshold and contour_h > light_height_threshold light_on_contours = [ contour_bounding_box_for_contour(c) for c in get_contours(light_on) if is_light_big_enough(c) ] light_off_contours = [ contour_bounding_box_for_contour(c) for c in get_contours(light_off) if is_light_big_enough(c) ] assert len(light_on_contours) + len(light_off_contours) == 1, \ "Expected to find exactly one light on the indicator" if light_on_contours: light_is_on = True light_contour = light_on_contours[0] else: light_is_on = False light_contour = light_off_contours[0] light_x, _ = get_center_for_contour(light_contour) if light_x > (w / 2.0): # Light is on the wrong side, need to flip 180 indicator = rotate_image_180(indicator) indicators_and_lights.append((indicator, light_is_on)) return indicators_and_lights
def get_connections(im): """ Returns a list of tuples, each with (color, destination, (to_click_x, to_click_y)). The list is ordered from top to bottom. """ blue = extract_color(im, 115, (100, 255), (0, 255)) red = extract_color(im, 5, (100, 255), (0, 255)) black = extract_color(im, (0, 255), (0, 255), (0, 5)) colors = ( (Color.BLUE, blue), (Color.BLACK, black), (Color.RED, red), ) mask = blue.copy() height, width = mask.shape rows = (0.33, 0.48, 0.62) cols = (0.30, 0.55) row_pxs = [int(row_percent * height) for row_percent in rows] start_col_px, end_col_px = [ int(col_percent * width) for col_percent in cols ] def _get_output_for_connection(color, start_row_idx, end_row_idx): """Returns a tuple with (color, destination, (to_click_x, to_click_y))""" # Use the start_row to find a point to click to cut the wire. to_click = (start_col_px, row_pxs[start_row_idx]) # Figure out which letter we're connecting to based on the end row. destination = ORDERED_DESTINATIONS[end_row_idx] return color, destination, to_click output = [] for start_row, start_row_px in enumerate(row_pxs): current_output = None for end_row, end_row_px in enumerate(row_pxs): # Reset the mask mask[:] = 0 # Draw a line where the connection should be on the mask cv2.line(mask, (start_col_px, start_row_px), (end_col_px, end_row_px), 255, 15) # Convert the mask to be usable in a masked_array bool_mask = np.invert(mask.astype(bool)) # We only want to consider it a wire there if there's at least 75% of the mask filled in activation_threshold = bool_mask.sum() * .66 for color, color_mat in colors: activated_amount = np.ma.masked_array(color_mat, mask=bool_mask).sum() if activated_amount > activation_threshold: new_output = _get_output_for_connection( color, start_row, end_row) new_output_and_activated = (new_output, activated_amount) if current_output is None: current_output = new_output_and_activated else: current_output = sorted( (current_output, new_output_and_activated), key=lambda x: x[1])[-1] # combined = np.zeros_like(im) # combined[:, :, 0] = color_mat # combined[:, :, 2] = mask # show(combined) if current_output is not None: print current_output output.append(current_output[0]) return output
def _get_wire_color_and_mat_or_none(wire, hue, saturation, value, color): mat = extract_color(wire, hue, saturation, value) if mat.any(): return color, mat else: return None
def _get_indicator_text(indicator, tesseract): indicator = get_subset(indicator, (40, 100), (0, 100)) color = extract_color(indicator, 47/2, (0, 50), (190, 240)) tesseract.SetImage(Image.fromarray(color)) return tesseract.GetUTF8Text().strip()
def get_is_done(im): im = get_subset(im, (80, 95), (5, 20)) color = extract_color(im, 128/2, (225, 255), (150, 255)) return cv2.countNonZero(color) != 0
def _get_cleaned_up_text_subsection(im): # type: (np.array) -> Optional[np.array] red1 = extract_color(im, (0, 6), (200, 255), (100, 150)) red2 = extract_color(im, (176, 180), (200, 255), (100, 150)) red = red1 + red2 color = extract_color(im, 45 / 2, (20, 50), (200, 255)) # show(red, .25) # show(yellow, .25) # im = scale(im, .25) # color = scale(color, .25) red_contour = _get_box_for_largest_rect_contour(red) text_contour = _get_box_for_largest_rect_contour(color) if red_contour is None or text_contour is None: # if red_contour is not None: # print "RED" # show(get_drawn_contours(red, [red_contour], True)) # show(color) # if text_contour is not None: # print "TEXT" # show(get_drawn_contours(color, [text_contour], True)) # show(red) # if not (red_contour is None and text_contour is None): # show(red) # show(color) # assert red_contour is None and text_contour is None, \ # "Error parsing serial number, didn't find one of the text or its label." return None red_center = get_center_for_contour(red_contour) text_center = get_center_for_contour(text_contour) text_subsection = four_point_transform(im, text_contour) # show(get_drawn_contours(color, text_contour, True), .25) # show(get_drawn_contours(red, red_contour, True), .25) # show(text_subsection) height, width = im.shape[:2] # Rotation logic from http://stackoverflow.com/a/5912847/3000133 if height > width: # Determine if red is left or right of text if text_center[0] < red_center[0]: text_subsection = rotate_image_counter_clockwise(text_subsection) else: # Rotate clockwise 90 text_subsection = rotate_image_clockwise(text_subsection) else: if text_center[1] > red_center[1]: # We're fine pass else: # Rotate 180 text_subsection = rotate_image_180(text_subsection) # show(get_drawn_contours(im, [text_contour], True)) # show(text_subsection) text_subsection_gray = cv2.cvtColor(text_subsection, cv2.COLOR_BGR2GRAY) # show(text_subsection_gray) _, text_threshold = cv2.threshold(text_subsection_gray, 50, 255, 0) text_threshold = 255 - text_threshold # show(text_threshold) height, width = text_threshold.shape[:2] text_threshold[:height / 10, :] = 0 text_threshold[9 * height / 10:, :] = 0 return text_threshold
def _get_indicator_text(indicator, tesseract): indicator = get_subset(indicator, (40, 100), (0, 100)) color = extract_color(indicator, 47 / 2, (0, 50), (190, 240)) tesseract.SetImage(Image.fromarray(color)) return tesseract.GetUTF8Text().strip()
def get_is_done(im): im = get_subset(im, (80, 95), (5, 20)) color = extract_color(im, 128 / 2, (225, 255), (150, 255)) return cv2.countNonZero(color) != 0