Spaces:
Sleeping
Sleeping
File size: 6,695 Bytes
f1554a2 5700698 f1554a2 5700698 b3e6fcf 5700698 f1554a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | 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
|