Spaces:
Sleeping
Sleeping
| import time | |
| import numpy as np | |
| import cv2 | |
| import pandas as pd | |
| from PIL import Image, ImageDraw, ImageFont | |
| class Element: | |
| def __init__(self, id, corner, category, text_content=None): | |
| self.id = id | |
| self.category = category | |
| self.col_min, self.row_min, self.col_max, self.row_max = corner | |
| self.width = self.col_max - self.col_min | |
| self.height = self.row_max - self.row_min | |
| self.area = self.width * self.height | |
| self.text_content = text_content | |
| self.parent_id = None | |
| self.children = [] # list of elements | |
| self.label = None | |
| def __str__(self): | |
| return (f"Element(id={self.id}, category={self.category}, corner=({self.col_min}, {self.row_min}, " | |
| f"{self.col_max}, {self.row_max}), width={self.width}, height={self.height}, area={self.area}, " | |
| f"text_content={self.text_content}, label={self.label}, parent_id={self.parent_id}, " | |
| f"children_count={len(self.children)})") | |
| def __eq__(self, other): | |
| if not isinstance(other, Element): | |
| return False | |
| return (self.id == other.id and | |
| self.category == other.category and | |
| self.col_min == other.col_min and | |
| self.row_min == other.row_min and | |
| self.col_max == other.col_max and | |
| self.row_max == other.row_max and | |
| self.text_content == other.text_content and | |
| self.label == other.label) | |
| def center(self): | |
| # 计算中心点 | |
| center_x = (self.col_min + self.col_max) / 2 | |
| center_y = (self.row_min + self.row_max) / 2 | |
| return (center_x, center_y) | |
| def init_bound(self): | |
| self.width = self.col_max - self.col_min | |
| self.height = self.row_max - self.row_min | |
| self.area = self.width * self.height | |
| def put_bbox(self): | |
| return self.col_min, self.row_min, self.col_max, self.row_max | |
| def wrap_info(self): | |
| info = {'id': self.id, 'class': self.category, 'height': self.height, 'width': self.width, | |
| 'position': {'column_min': self.col_min, 'row_min': self.row_min, 'column_max': self.col_max, | |
| 'row_max': self.row_max}, 'label': self.label} | |
| if self.text_content is not None: | |
| info['text_content'] = self.text_content | |
| if len(self.children) > 0: | |
| info['children'] = [] | |
| for child in self.children: | |
| info['children'].append(child.id) | |
| if self.parent_id is not None: | |
| info['parent'] = self.parent_id | |
| return info | |
| def resize(self, resize_ratio): | |
| self.col_min = int(self.col_min * resize_ratio) | |
| self.row_min = int(self.row_min * resize_ratio) | |
| self.col_max = int(self.col_max * resize_ratio) | |
| self.row_max = int(self.row_max * resize_ratio) | |
| self.init_bound() | |
| def element_merge(self, element_b, new_element=False, new_category=None, new_id=None): | |
| col_min_a, row_min_a, col_max_a, row_max_a = self.put_bbox() | |
| col_min_b, row_min_b, col_max_b, row_max_b = element_b.put_bbox() | |
| new_corner = ( | |
| min(col_min_a, col_min_b), min(row_min_a, row_min_b), max(col_max_a, col_max_b), max(row_max_a, row_max_b)) | |
| if element_b.text_content is not None: | |
| self.text_content = element_b.text_content if self.text_content is None else self.text_content + '\n' + element_b.text_content | |
| if new_element: | |
| return Element(new_id, new_corner, new_category) | |
| else: | |
| self.col_min, self.row_min, self.col_max, self.row_max = new_corner | |
| self.init_bound() | |
| def calc_intersection_area(self, element_b, bias=(0, 0)): | |
| a = self.put_bbox() | |
| b = element_b.put_bbox() | |
| col_min_s = max(a[0], b[0]) - bias[0] | |
| row_min_s = max(a[1], b[1]) - bias[1] | |
| col_max_s = min(a[2], b[2]) | |
| row_max_s = min(a[3], b[3]) | |
| w = np.maximum(0, col_max_s - col_min_s) | |
| h = np.maximum(0, row_max_s - row_min_s) | |
| inter = w * h | |
| iou = inter / (self.area + element_b.area - inter) | |
| ioa = inter / self.area | |
| iob = inter / element_b.area | |
| return inter, iou, ioa, iob | |
| def element_relation(self, element_b, bias=(0, 0)): | |
| """ | |
| @bias: (horizontal bias, vertical bias) | |
| :return: -1 : a in b | |
| 0 : a, b are not intersected | |
| 1 : b in a | |
| 2 : a, b are identical or intersected | |
| """ | |
| inter, iou, ioa, iob = self.calc_intersection_area(element_b, bias) | |
| # area of intersection is 0 | |
| if ioa == 0: | |
| return 0 | |
| # a in b | |
| if ioa >= 1: | |
| return -1 | |
| # b in a | |
| if iob >= 1: | |
| return 1 | |
| return 2 | |
| # def visualize_element(self, img, color=(0, 255, 0), line=3, show=False, ratio=1, expand_size=5, corner_radius=20): | |
| # loc = self.put_bbox() | |
| # | |
| # if ratio != 1: | |
| # loc = [int(x * ratio) for x in loc] | |
| # | |
| # # 调整线条粗细,根据缩放比例 | |
| # adjusted_thickness = int(line * ratio) | |
| # if adjusted_thickness < 1: | |
| # adjusted_thickness = 1 # 确保最小粗细为1 | |
| # | |
| # # cv2.rectangle(img, loc[:2], loc[2:], color, line) | |
| # # cv2.rectangle(img, (loc[0], loc[1]), (loc[2], loc[3]), color, line) | |
| # # 标号 | |
| # # cv2.putText(img, str(int(self.id) + 1), (int(ratio*(self.col_min - 10)), int(ratio*(self.row_max + 10))), | |
| # # cv2.FONT_HERSHEY_SIMPLEX, 1, color, line) | |
| # | |
| # # 扩展边框,在原始坐标的基础上增加 expand_size | |
| # loc[0] -= expand_size # 左 | |
| # loc[1] -= expand_size + 2 # 上 | |
| # loc[2] += expand_size + 10 # 右 | |
| # loc[3] += expand_size + 2 # 下 | |
| # | |
| # # 确保圆角半径不会超过矩形的宽或高的一半 | |
| # corner_radius = min(corner_radius, (loc[2] - loc[0]) // 2, (loc[3] - loc[1]) // 2) | |
| # | |
| # # 绘制圆角矩形 | |
| # def draw_rounded_rectangle(image, top_left, bottom_right, color, thickness, radius): | |
| # x1, y1 = top_left | |
| # x2, y2 = bottom_right | |
| # | |
| # # 计算各条边线的起点和终点 | |
| # points = [ | |
| # ((x1 + radius, y1), (x2 - radius, y1)), # 上边 | |
| # ((x1 + radius, y2), (x2 - radius, y2)), # 下边 | |
| # ((x1, y1 + radius), (x1, y2 - radius)), # 左边 | |
| # ((x2, y1 + radius), (x2, y2 - radius)) # 右边 | |
| # ] | |
| # | |
| # # 画四分之一圆角 | |
| # cv2.ellipse(image, (x1 + radius, y1 + radius), (radius, radius), 180, 0, 90, color, thickness) # 左上角 | |
| # cv2.ellipse(image, (x2 - radius, y1 + radius), (radius, radius), 270, 0, 90, color, thickness) # 右上角 | |
| # cv2.ellipse(image, (x1 + radius, y2 - radius), (radius, radius), 90, 0, 90, color, thickness) # 左下角 | |
| # cv2.ellipse(image, (x2 - radius, y2 - radius), (radius, radius), 0, 0, 90, color, thickness) # 右下角 | |
| # | |
| # # 画四条直线连接四分之一圆角 | |
| # for point_start, point_end in points: | |
| # cv2.line(image, point_start, point_end, color, thickness) | |
| # | |
| # # 绘制圆角矩形 | |
| # draw_rounded_rectangle(img, (loc[0], loc[1]), (loc[2], loc[3]), color, adjusted_thickness, corner_radius) | |
| # | |
| # # for child in self.children: | |
| # # child.visualize_element(img, color=(255, 0, 255), line=line) | |
| # if show: | |
| # cv2.imshow('element', img) | |
| # cv2.waitKey(0) | |
| # cv2.destroyWindow('element') | |
| def adjust_loc(self, ratio=1, expand_size=5): | |
| loc = list(self.put_bbox()) | |
| if ratio != 1: | |
| loc = [int(x * ratio) for x in loc] | |
| # 扩展框的边界 | |
| loc[0] -= expand_size # 左扩展 | |
| loc[1] -= expand_size + 2 # 上扩展 | |
| loc[2] += expand_size + 10 # 右扩展 | |
| loc[3] += expand_size + 2 # 下扩展 | |
| return loc | |
| def draw_and_expand_rounded_rectangle(self, img, loc, color=(0, 255, 0), thickness=-1, expand_size=5, | |
| corner_radius=20, ): | |
| """ | |
| 通用的绘制带扩展边框的圆角矩形的方法。 | |
| 参数: | |
| - img: 输入图像 | |
| - loc: 矩形的左上角和右下角坐标 (x1, y1, x2, y2) | |
| - color: 矩形颜色 | |
| - line: 矩形边框线条的粗细 | |
| - expand_size: 扩展的边框大小 | |
| - corner_radius: 圆角的半径 | |
| - thickness: 线条粗细(-1 为填充) | |
| """ | |
| # 确保圆角半径不会超过矩形的宽或高的一半 | |
| corner_radius = min(corner_radius, (loc[2] - loc[0]) // 2, (loc[3] - loc[1]) // 2) | |
| # 绘制带扩展的圆角矩形 | |
| x1, y1 = loc[0], loc[1] | |
| x2, y2 = loc[2], loc[3] | |
| x1 = int(x1) | |
| y1 = int(y1) | |
| x2 = int(x2) | |
| y2 = int(y2) | |
| # Step 1: 处理填充 | |
| if thickness == -1: | |
| # 填充模式 - 先绘制没有圆角的矩形部分 | |
| cv2.rectangle(img, (x1 + corner_radius, y1), (x2 - corner_radius, y2), color, cv2.FILLED) # 上下边部分 | |
| cv2.rectangle(img, (x1, y1 + corner_radius), (x2, y2 - corner_radius), color, cv2.FILLED) # 左右边部分 | |
| # Step 2: 处理圆角 | |
| # 画四分之一圆角 | |
| cv2.ellipse(img, (x1 + corner_radius, y1 + corner_radius), (corner_radius, corner_radius), 180, 0, 90, color, | |
| thickness) # 左上角 | |
| cv2.ellipse(img, (x2 - corner_radius, y1 + corner_radius), (corner_radius, corner_radius), 270, 0, 90, color, | |
| thickness) # 右上角 | |
| cv2.ellipse(img, (x1 + corner_radius, y2 - corner_radius), (corner_radius, corner_radius), 90, 0, 90, color, | |
| thickness) # 左下角 | |
| cv2.ellipse(img, (x2 - corner_radius, y2 - corner_radius), (corner_radius, corner_radius), 0, 0, 90, color, | |
| thickness) # 右下角 | |
| # Step 3: 处理边框 | |
| if thickness > 0: | |
| # 绘制矩形四条边框,连接四个圆角 | |
| cv2.line(img, (x1 + corner_radius, y1), (x2 - corner_radius, y1), color, thickness) # 上边 | |
| cv2.line(img, (x1 + corner_radius, y2), (x2 - corner_radius, y2), color, thickness) # 下边 | |
| cv2.line(img, (x1, y1 + corner_radius), (x1, y2 - corner_radius), color, thickness) # 左边 | |
| cv2.line(img, (x2, y1 + corner_radius), (x2, y2 - corner_radius), color, thickness) # 右边 | |
| def visualize_element(self, img, color=(0, 255, 0), line=3, show=False, ratio=1, expand_size=5, corner_radius=20, | |
| loc=None): | |
| # 调整线条粗细,根据缩放比例 | |
| loc = self.adjust_loc(ratio=ratio, expand_size=expand_size) | |
| if loc is None: | |
| loc = [0, 0, 0, 0] | |
| adjusted_thickness = int(line * ratio) | |
| if adjusted_thickness < 1: | |
| adjusted_thickness = 1 # 确保最小粗细为1 | |
| # 调用通用的绘制带扩展的圆角矩形方法 | |
| self.draw_and_expand_rounded_rectangle(img, loc, color, adjusted_thickness, expand_size, corner_radius) | |
| # 显示图像(可选) | |
| if show: | |
| cv2.imshow('element', img) | |
| cv2.waitKey(0) | |
| cv2.destroyWindow('element') | |
| def calculate_text_height(self, font_title, font_text, title, text, bubble_width, inner_padding): | |
| """ | |
| 计算文本内容的高度,包括标题和正文 | |
| """ | |
| # 创建临时 PIL 图像用于计算文本高度 | |
| temp_img = Image.new("RGB", (bubble_width, 500), (255, 255, 255)) | |
| draw = ImageDraw.Draw(temp_img) | |
| # 计算标题的高度 | |
| title_height = draw.textbbox((0, 0), title, font=font_title)[3] - \ | |
| draw.textbbox((0, 0), title, font=font_title)[1] | |
| # 处理正文的自动换行 | |
| words = text.split(' ') | |
| current_line = "" | |
| lines = [] | |
| max_text_width = bubble_width - 2 * inner_padding | |
| for word in words: | |
| test_line = current_line + word + " " | |
| text_bbox = draw.textbbox((0, 0), test_line, font=font_text) | |
| text_width = text_bbox[2] - text_bbox[0] | |
| if text_width < max_text_width: | |
| current_line = test_line | |
| else: | |
| lines.append(current_line) | |
| current_line = word + " " | |
| # 添加最后一行 | |
| lines.append(current_line) | |
| # 计算正文高度 | |
| text_height = len(lines) * (draw.textbbox((0, 0), lines[0], font=font_text)[3] - | |
| draw.textbbox((0, 0), lines[0], font=font_text)[1]) + len(lines) * 5 | |
| # 计算总高度(标题 + 正文 + 内边距) | |
| total_height = title_height + text_height + 2 * inner_padding + 50 | |
| return total_height | |
| def write_text(self, img, bubble_top_left, font_title, font_text, title, text, bubble_width, inner_padding): | |
| """ | |
| 在对话框内绘制标题和正文内容 | |
| """ | |
| # 确保文本和对话框边缘有内边距 | |
| text_area_top_left = (bubble_top_left[0] + inner_padding, bubble_top_left[1] + inner_padding - 15) | |
| # 将 OpenCV 图像转换为 PIL 图像 | |
| img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) | |
| # 创建 ImageDraw 对象用于绘制文字 | |
| draw = ImageDraw.Draw(img_pil) | |
| # 绘制标题 | |
| draw.text((text_area_top_left[0], text_area_top_left[1]), title, font=font_title, fill=(255, 255, 255)) | |
| # 处理正文自动换行 | |
| text_start_y = text_area_top_left[1] + font_title.size + 50 # 标题下面再留些空间给正文 | |
| max_text_width = bubble_width - 2 * inner_padding | |
| lines = [] | |
| words = text.split(' ') | |
| current_line = "" | |
| for word in words: | |
| # 检查当前行的宽度 | |
| test_line = current_line + word + " " | |
| text_bbox = draw.textbbox((0, 0), test_line, font=font_text) | |
| text_width = text_bbox[2] - text_bbox[0] | |
| if text_width < max_text_width: | |
| current_line = test_line | |
| else: | |
| # 当前行内容超出对话框宽度,换行 | |
| lines.append(current_line) | |
| current_line = word + " " | |
| # 添加最后一行 | |
| lines.append(current_line) | |
| # 绘制每一行正文 | |
| for line in lines: | |
| draw.text((text_area_top_left[0], text_start_y), line, font=font_text, fill=(255, 255, 255)) | |
| text_start_y += font_text.size + 5 | |
| # 将 PIL 图像转换回 OpenCV 图像(直接修改 img) | |
| img[:] = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR) | |
| def process_img(self, img, elements, texts, color=(0, 255, 0), line=3, alpha=0.7, expand_size=5, | |
| corner_radius=20, show=False, ratio=1, padding=50, triangle_height=0, inner_padding=100, | |
| font_path_title="CDM/detect_merge/title.ttf", font_path_text="CDM/detect_merge/text.ttf", | |
| font_size_title=70, font_size_text=50): | |
| # --------------画所有边框+黑色半透明处理剩余部分---------------------- | |
| black_block_start = time.time() | |
| # 创建一个全黑遮罩(用于透明处理) | |
| mask = np.zeros_like(img, dtype=np.uint8) | |
| # 遍历每个元素并绘制圆角矩形,排除这些区域 | |
| for element in elements: | |
| loc = element.adjust_loc(ratio=ratio, expand_size=expand_size) | |
| # 调用通用的绘制带扩展的圆角矩形方法,在遮罩上扣除这些区域 | |
| self.draw_and_expand_rounded_rectangle(mask, loc, (255, 255, 255), -1, expand_size, corner_radius) | |
| # 创建一个全黑的图像,用于制作黑色遮罩 | |
| overlay = np.zeros_like(img, dtype=np.uint8) | |
| overlay[:] = (0, 0, 0) # 这是一个全黑的遮罩 | |
| # 将 overlay 和 img 混合,但根据 mask 的黑色区域应用透明效果,白色区域保持原样 | |
| img_with_overlay = cv2.addWeighted(img, 1 - alpha, overlay, alpha, 0) # 将黑色透明度应用到整个图像 | |
| # 使用 mask 保留元素区域的原始图像 | |
| result = cv2.bitwise_and(img, mask) # 仅保留元素区域(白色部分) | |
| result += cv2.bitwise_and(img_with_overlay, cv2.bitwise_not(mask)) # 将黑色透明遮罩应用到非元素区域(黑色部分) | |
| # 将结果应用到 img 上 | |
| img[:, :, :] = result | |
| # 为每个元素再绘制圆角矩形边框 | |
| for element in elements: | |
| loc = element.adjust_loc(ratio=ratio, expand_size=expand_size) | |
| element.visualize_element(img, color=color, line=line, show=show, ratio=ratio, loc=loc) | |
| black_block_cost = time.time() - black_block_start | |
| # print("绘制带黑色遮罩的框选component的圆角block花费:%2.2f s" % black_block_cost) | |
| # -----------------------绘制信息部分---------------------------------- | |
| write_text_in_block = time.time() | |
| # 获取图像的宽度和高度 | |
| img_height, img_width = img.shape[:2] | |
| # 获取标题和正文的显示信息 | |
| title = "· " + elements[0].label | |
| df = pd.DataFrame(texts) | |
| df_unique = df.drop_duplicates(subset='label', keep='first') | |
| text = df_unique[df_unique['label'] == elements[0].label]['segment'].values[0] | |
| text = text[0].upper() + text[1:] | |
| # 加载字体 | |
| font_title = ImageFont.truetype(font_path_title, font_size_title) | |
| font_text = ImageFont.truetype(font_path_text, font_size_text) | |
| # 计算所有元素的调整后位置 | |
| locs = [e.adjust_loc(ratio=ratio, expand_size=expand_size) for e in elements] | |
| # 获取所有元素的上边界(y1)和下边界(y2) | |
| # 并添加图片的上边界(0)和下边界(img_height)作为特殊的"元素" | |
| y_positions = [] | |
| for loc in locs: | |
| y_positions.append((loc[1], 'element_top')) # 元素的上边界 | |
| y_positions.append((loc[3], 'element_bottom')) # 元素的下边界 | |
| # 添加图片的上边界和下边界 | |
| y_positions.append((0, 'image_top')) | |
| y_positions.append((img_height, 'image_bottom')) | |
| # 按 Y 坐标排序 | |
| y_positions.sort(key=lambda x: x[0]) | |
| # 找到所有的垂直空白区域 | |
| gaps = [] | |
| for i in range(len(y_positions) - 1): | |
| y1, label1 = y_positions[i] | |
| y2, label2 = y_positions[i + 1] | |
| # 如果两个位置之间有空隙 | |
| if y2 > y1: | |
| gap_height = y2 - y1 | |
| gap = { | |
| 'start': y1, | |
| 'end': y2, | |
| 'height': gap_height, | |
| 'above': label1, | |
| 'below': label2 | |
| } | |
| gaps.append(gap) | |
| # 找到高度最大的空白区域 | |
| largest_gap = max(gaps, key=lambda x: x['height']) | |
| # 判断最大的空白区域是否在元素和图片边界之间,还是在元素之间 | |
| is_between_element_and_border = False | |
| is_between_elements = False | |
| if largest_gap['above'] == 'image_top' or largest_gap['below'] == 'image_bottom': | |
| is_between_element_and_border = True | |
| elif largest_gap['above'] == 'element_bottom' and largest_gap['below'] == 'element_top': | |
| is_between_elements = True | |
| # 计算文本框的宽度 | |
| bubble_width = img_width - 2 * padding | |
| # 计算文本框的高度 | |
| bubble_height = self.calculate_text_height(font_title, font_text, title, text, bubble_width, inner_padding) | |
| # 确保文本框的高度不超过最大空白区域的高度减去两倍的内边距 | |
| max_bubble_height = largest_gap['height'] - 2 * padding | |
| if bubble_height > max_bubble_height: | |
| bubble_height = max_bubble_height | |
| # 根据判断结果确定文本框的位置 | |
| if is_between_element_and_border: | |
| # 文本框靠近对应的元素 | |
| if largest_gap['above'] == 'image_top': | |
| # 空白区域在顶部,文本框在第一个元素的上方 | |
| ref_element = min(elements, key=lambda e: e.adjust_loc(ratio=ratio, expand_size=expand_size)[1]) | |
| loc = ref_element.adjust_loc(ratio=ratio, expand_size=expand_size) | |
| y1 = loc[1] | |
| bubble_top = y1 - triangle_height - padding - bubble_height | |
| bubble_left = padding | |
| elif largest_gap['below'] == 'image_bottom': | |
| # 空白区域在底部,文本框在最后一个元素的下方 | |
| ref_element = max(elements, key=lambda e: e.adjust_loc(ratio=ratio, expand_size=expand_size)[3]) | |
| loc = ref_element.adjust_loc(ratio=ratio, expand_size=expand_size) | |
| y2 = loc[3] | |
| bubble_top = y2 + triangle_height + padding | |
| bubble_left = padding | |
| bubble_top_left = (bubble_left, bubble_top) | |
| bubble_bottom_right = (bubble_left + bubble_width, bubble_top + bubble_height) | |
| elif is_between_elements: | |
| # 文本框在元素之间,垂直居中放置在最大空白区域内 | |
| gap_start = largest_gap['start'] | |
| gap_end = largest_gap['end'] | |
| bubble_top = gap_start + (gap_end - gap_start - bubble_height) / 2 | |
| bubble_left = padding | |
| bubble_top_left = (bubble_left, bubble_top) | |
| bubble_bottom_right = (bubble_left + bubble_width, bubble_top + bubble_height) | |
| else: | |
| # 如果没有符合条件的空白区域,默认将文本框放在图片的顶部 | |
| bubble_top_left = (padding, padding) | |
| bubble_bottom_right = (bubble_top_left[0] + bubble_width, bubble_top_left[1] + bubble_height) | |
| # 确保文本框不会超出图片的顶部或底部边界 | |
| if bubble_top_left[1] < 0: | |
| bubble_top_left = (bubble_top_left[0], padding) | |
| bubble_bottom_right = (bubble_bottom_right[0], bubble_top_left[1] + bubble_height) | |
| elif bubble_bottom_right[1] > img_height: | |
| bubble_bottom_right = (bubble_bottom_right[0], img_height - padding) | |
| bubble_top_left = (bubble_top_left[0], bubble_bottom_right[1] - bubble_height) | |
| # 绘制圆角矩形作为对话框 | |
| self.draw_and_expand_rounded_rectangle( | |
| img, | |
| [bubble_top_left[0], bubble_top_left[1], bubble_bottom_right[0], bubble_bottom_right[1]], | |
| color, | |
| corner_radius=50 | |
| ) | |
| # 对话框内写入文本 | |
| self.write_text(img, bubble_top_left, font_title, font_text, title, text, bubble_width, inner_padding) | |
| # # 计算所有元素的调整后位置 | |
| # locs = [e.adjust_loc(ratio=ratio, expand_size=expand_size) for e in elements] | |
| # | |
| # # 计算元素组的包围框 | |
| # min_y1 = min(loc[1] for loc in locs) | |
| # max_y2 = max(loc[3] for loc in locs) | |
| # | |
| # # 计算元素组到图片上下边界的距离 | |
| # top_distance_group = min_y1 # 元素组上边到图片上边界的距离 | |
| # bottom_distance_group = img_height - max_y2 # 元素组下边到图片下边界的距离 | |
| # | |
| # if top_distance_group > bottom_distance_group: | |
| # # 如果元素组上方空间更大,选择最靠近上边界的元素 | |
| # ref_element = min(elements, key=lambda e: e.adjust_loc(ratio=ratio, expand_size=expand_size)[1]) | |
| # else: | |
| # # 否则,选择最靠近下边界的元素 | |
| # ref_element = max(elements, key=lambda e: e.adjust_loc(ratio=ratio, expand_size=expand_size)[3]) | |
| # | |
| # # 使用参考元素进行后续操作 | |
| # loc = ref_element.adjust_loc(ratio=ratio, expand_size=expand_size) | |
| # # 计算每个扩展后边界到图片边界的距离 | |
| # left_distance = loc[0] # 左边到图像左边界的距离 | |
| # right_distance = img_width - loc[2] # 右边到图像右边界的距离 | |
| # top_distance = loc[1] # 上边到图像上边界的距离 | |
| # bottom_distance = img_height - loc[3] # 下边到图像下边界的距离 | |
| # | |
| # # 存储每个边界的距离 | |
| # distances = { | |
| # 'left': left_distance, | |
| # 'right': right_distance, | |
| # 'top': top_distance, | |
| # 'bottom': bottom_distance | |
| # } | |
| # | |
| # # 计算对话框宽度(接近 img 的宽度,但稍微小于 img 宽度) | |
| # bubble_width = img_width - 2 * padding | |
| # x1, y1, x2, y2 = loc | |
| # | |
| # # 计算文本框的高度 | |
| # bubble_height = self.calculate_text_height(font_title, font_text, title, text, bubble_width, inner_padding) | |
| # | |
| # # 根据元素位置绘制对话框 | |
| # if distances['top'] > distances['bottom']: | |
| # # 对话框的位置:在元素上方 | |
| # bubble_top_left = (padding, y1 - triangle_height - padding - bubble_height) | |
| # bubble_bottom_right = (bubble_width + padding, y1 - triangle_height - padding) | |
| # else: | |
| # # 对话框的位置:在元素下方 | |
| # bubble_top_left = (padding, y2 + triangle_height + padding) | |
| # bubble_bottom_right = (bubble_width + padding, y2 + triangle_height + padding + bubble_height) | |
| # | |
| # # 绘制圆角矩形作为对话框 | |
| # self.draw_and_expand_rounded_rectangle( | |
| # img, | |
| # [bubble_top_left[0], bubble_top_left[1], bubble_bottom_right[0], bubble_bottom_right[1]], | |
| # color, | |
| # corner_radius=50 | |
| # ) | |
| # | |
| # # 对话框内写入文本 | |
| # self.write_text(img, bubble_top_left, font_title, font_text, title, text, bubble_width, inner_padding) | |
| write_in_block_cost = time.time() - write_text_in_block | |
| # print("绘制文本展示block+写入文字花费时间:%2.2f s" % write_in_block_cost) | |
| # ===================================================== | |
| # # 如果这张图只有一个element | |
| # if len(elements) == 1: | |
| # loc = elements[0].adjust_loc(ratio=ratio, expand_size=expand_size) | |
| # # 计算每个扩展后边界到图片边界的距离 | |
| # left_distance = loc[0] # 左边到图像左边界的距离 | |
| # right_distance = img_width - loc[2] # 右边到图像右边界的距离 | |
| # top_distance = loc[1] # 上边到图像上边界的距离 | |
| # bottom_distance = img_height - loc[3] # 下边到图像下边界的距离 | |
| # | |
| # # 存储每个边界的距离 | |
| # distances = { | |
| # 'left': left_distance, | |
| # 'right': right_distance, | |
| # 'top': top_distance, | |
| # 'bottom': bottom_distance | |
| # } | |
| # | |
| # # 计算对话框宽度(接近 img 的宽度,但稍微小于 img 宽度) | |
| # bubble_width = img_width - 2 * padding | |
| # x1, y1, x2, y2 = loc | |
| # | |
| # # 计算文本框的高度 | |
| # bubble_height = self.calculate_text_height(font_title, font_text, title, text, bubble_width, inner_padding) | |
| # | |
| # # 根据元素位置绘制对话框 | |
| # if distances['top'] > distances['bottom']: | |
| # # 对话框的位置:在元素上方 | |
| # bubble_top_left = (padding, y1 - triangle_height - padding - bubble_height) | |
| # bubble_bottom_right = (bubble_width + padding, y1 - triangle_height - padding) | |
| # # # 三角形的底边中心位置 (等边三角形) | |
| # # triangle_base_y = bubble_bottom_right[1] | |
| # # triangle_center_x = (x1 + x2) // 2 | |
| # # # 三角形顶点朝下 | |
| # # triangle_tip_y = triangle_base_y + (triangle_height - 10) | |
| # else: | |
| # # 对话框的位置:在元素下方 | |
| # bubble_top_left = (padding, y2 + triangle_height + padding) | |
| # bubble_bottom_right = (bubble_width + padding, y2 + triangle_height + padding + bubble_height) | |
| # # # 三角形的底边中心位置 (等边三角形) | |
| # # triangle_base_y = bubble_top_left[1] | |
| # # triangle_center_x = (x1 + x2) // 2 | |
| # # # 三角形顶点朝上 | |
| # # triangle_tip_y = triangle_base_y - (triangle_height - 10) | |
| # | |
| # # # 三角形底边长度等于高度,所以底边是 [triangle_height] 的两倍 | |
| # # half_base = triangle_height // 2 + 5 | |
| # # | |
| # # # 绘制等边三角形 | |
| # # triangle_points = np.array([ | |
| # # [triangle_center_x - half_base, triangle_base_y], # 三角形左边点 | |
| # # [triangle_center_x + half_base, triangle_base_y], # 三角形右边点 | |
| # # [triangle_center_x, triangle_tip_y] # 三角形顶点 | |
| # # ]) | |
| # # cv2.fillPoly(img, [triangle_points], color) | |
| # | |
| # # 绘制圆角矩形作为对话框 | |
| # self.draw_and_expand_rounded_rectangle(img, | |
| # [bubble_top_left[0], bubble_top_left[1], bubble_bottom_right[0], | |
| # bubble_bottom_right[1]], color, corner_radius=50) | |
| # | |
| # # 对话框内写入文本 | |
| # self.write_text(img, bubble_top_left, font_title, font_text, title, text, bubble_width, inner_padding) | |
| # | |
| # # else: | |