Spaces:
Sleeping
Sleeping
| import cv2 | |
| import numpy as np | |
| from random import randint as rint | |
| from CDM.config.CONFIG_UIED import Config | |
| from os.path import join as pjoin | |
| C = Config() | |
| def draw_bounding_box_class_by_category(org, components, color_map=C.COLOR, line=2, show=False, output_root=None): | |
| """ | |
| Draw bounding box of components and save images based on their classes. | |
| Each class of components will be saved in a different image. | |
| :param org: original image | |
| :param components: bbox [(column_min, row_min, column_max, row_max)] -> components list | |
| :param color_map: colors mapping to different components | |
| :param line: line thickness | |
| :param show: show or not | |
| :param output_root: the directory to save classified images | |
| """ | |
| # Create a copy of the original image for each category | |
| boards_by_category = {} | |
| for compo in components: | |
| # Get the bounding box of the component | |
| bbox = compo.put_bbox() | |
| category = compo.category # Get the category of the component | |
| # Initialize a board for each category if not already done | |
| if category not in boards_by_category: | |
| boards_by_category[category] = org.copy() | |
| # Draw the bounding box on the board of the corresponding category | |
| boards_by_category[category] = cv2.rectangle( | |
| boards_by_category[category], | |
| (bbox[0], bbox[1]), (bbox[2], bbox[3]), | |
| color_map[category], line | |
| ) | |
| # Save or show the results for each category | |
| for category, board in boards_by_category.items(): | |
| if show: | |
| cv2.imshow(category, board) | |
| cv2.waitKey(0) | |
| if output_root is not None: | |
| # Create the output path for this category | |
| output_path = pjoin(output_root, f"{category}.jpg") | |
| cv2.imwrite(output_path, board) | |
| # 返回所有分类对应的图像字典 | |
| return boards_by_category | |
| 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, ratio, 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() | |
| # board = cv2.imread(img_path) | |
| # ratio = board.shape[0]/org.shape[0] | |
| for compo in components: | |
| bbox = compo.put_bbox() | |
| # bounding box on full size image | |
| # bbox = int(ratio * bbox) | |
| bbox = [int(x * ratio) for x in 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 | |