import cv2 import numpy as np from random import randint as rint from config.CONFIG_UIED import Config C = Config() def draw_bounding_box_class(org, components, color_map=C.COLOR, line=2, show=False, write_path=None, name='board'): """ Draw bounding box of components with their classes on the original image :param org: original image :param components: bbox [(column_min, row_min, column_max, row_max)] -> top_left: (column_min, row_min) -> bottom_right: (column_max, row_max) :param color_map: colors mapping to different components :param line: line thickness :param compo_class: classes matching the corners of components :param show: show or not :return: labeled image """ board = org.copy() for compo in components: bbox = compo.put_bbox() board = cv2.rectangle(board, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color_map[compo.category], line) # board = cv2.putText(board, compo.category, (bbox[0]+5, bbox[1]+20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color_map[compo.category], 2) if show: cv2.imshow(name, board) cv2.waitKey(0) if write_path is not None: cv2.imwrite(write_path, board) return board def draw_bounding_box(org, components, color=(0, 255, 0), line=2, show=False, write_path=None, name='board', is_return=False, wait_key=0): """ Draw bounding box of components on the original image :param org: original image :param components: bbox [(column_min, row_min, column_max, row_max)] -> top_left: (column_min, row_min) -> bottom_right: (column_max, row_max) :param color: line color :param line: line thickness :param show: show or not :return: labeled image """ if not show and write_path is None and not is_return: return board = org.copy() for compo in components: bbox = compo.put_bbox() board = cv2.rectangle(board, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, line) if show: cv2.imshow(name, board) if wait_key is not None: cv2.waitKey(wait_key) if wait_key == 0: cv2.destroyWindow(name) if write_path is not None: # board = cv2.resize(board, (1080, 1920)) # board = board[100:-110] cv2.imwrite(write_path, board) return board def draw_line(org, lines, color=(0, 255, 0), show=False): """ Draw detected lines on the original image :param org: original image :param lines: [line_h, line_v] -> line_h: horizontal {'head':(column_min, row), 'end':(column_max, row), 'thickness':int) -> line_v: vertical {'head':(column, row_min), 'end':(column, row_max), 'thickness':int} :param color: drawn color :param show: show or not :return: image with lines drawn """ board = org.copy() line_h, line_v = lines for line in line_h: cv2.line(board, tuple(line['head']), tuple(line['end']), color, line['thickness']) for line in line_v: cv2.line(board, tuple(line['head']), tuple(line['end']), color, line['thickness']) if show: cv2.imshow('img', board) cv2.waitKey(0) return board def draw_boundary(components, shape, show=False): """ Draw boundary of objects on the black withe :param components: boundary: [top, bottom, left, right] -> up, bottom: (column_index, min/max row border) -> left, right: (row_index, min/max column border) detect range of each row :param shape: shape or original image :param show: show or not :return: drawn board """ board = np.zeros(shape[:2], dtype=np.uint8) # binary board for component in components: # up and bottom: (column_index, min/max row border) for point in component.boundary[0] + component.boundary[1]: board[point[1], point[0]] = 255 # left, right: (row_index, min/max column border) for point in component.boundary[2] + component.boundary[3]: board[point[0], point[1]] = 255 if show: cv2.imshow('rec', board) cv2.waitKey(0) return board def draw_region(region, broad, show=False): color = (rint(0,255), rint(0,255), rint(0,255)) for point in region: broad[point[0], point[1]] = color if show: cv2.imshow('region', broad) cv2.waitKey() return broad def draw_region_bin(region, broad, show=False): for point in region: broad[point[0], point[1]] = 255 if show: cv2.imshow('region', broad) cv2.waitKey() return broad