733 lines
25 KiB
Python
733 lines
25 KiB
Python
import base64
|
||
import io
|
||
import os
|
||
from typing import Any, Dict, List, Literal, Optional, Tuple, Union
|
||
|
||
# 在任何 matplotlib 导入之前强制设置非交互式后端
|
||
# 避免在 macOS 自动化版本中弹出 Quartz 窗口导致阴塞
|
||
import matplotlib
|
||
matplotlib.use('Agg')
|
||
|
||
import cv2
|
||
import easyocr
|
||
import numpy as np
|
||
import torch
|
||
from PIL import Image
|
||
from PIL.Image import Image as ImageType
|
||
from supervision.detection.core import Detections
|
||
from supervision.draw.color import Color, ColorPalette
|
||
from torchvision.ops import box_convert
|
||
from torchvision.transforms import ToPILImage
|
||
from transformers import AutoModelForCausalLM, AutoProcessor
|
||
from transformers.image_utils import load_image
|
||
from ultralytics import YOLO
|
||
|
||
class EndpointHandler:
|
||
"""
|
||
OmniParser 推理处理器
|
||
|
||
Args:
|
||
model_dir: 模型目录路径
|
||
enable_ocr: 是否启用 OCR 文字识别（默认 True）
|
||
enable_caption: 是否启用图标描述生成（默认 True）
|
||
enable_yolo: 是否启用 YOLO 图标检测（默认 True）
|
||
ocr_languages: OCR 语言列表（默认 ["ch_sim", "en"] 中英文）
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
model_dir: str = os.path.dirname(os.path.abspath(__file__)),
|
||
enable_ocr: bool = True,
|
||
enable_caption: bool = True,
|
||
enable_yolo: bool = True,
|
||
ocr_languages: Optional[List[str]] = None,
|
||
) -> None:
|
||
# 保存开关状态
|
||
self.enable_ocr = enable_ocr
|
||
self.enable_caption = enable_caption
|
||
self.enable_yolo = enable_yolo
|
||
|
||
# 选择计算设备
|
||
if torch.cuda.is_available():
|
||
self.device = torch.device("cuda")
|
||
device_name = torch.cuda.get_device_name(0)
|
||
print(f"🚀 使用设备: GPU ({device_name})")
|
||
elif torch.backends.mps.is_available():
|
||
self.device = torch.device("mps")
|
||
print("🚀 使用设备: GPU (Apple MPS)")
|
||
else:
|
||
self.device = torch.device("cpu")
|
||
print("⚠️ 使用设备: CPU (推理较慢，建议安装 CUDA 版 PyTorch)")
|
||
|
||
# 显示启用的模块
|
||
modules = []
|
||
if enable_yolo:
|
||
modules.append("YOLO检测")
|
||
if enable_ocr:
|
||
modules.append("OCR文字")
|
||
if enable_caption:
|
||
modules.append("图标描述")
|
||
print(f"📦 启用模块: {', '.join(modules) if modules else '无'}")
|
||
|
||
# YOLO 图标检测模型
|
||
self.yolo = None
|
||
if enable_yolo:
|
||
self.yolo = YOLO(f"{model_dir}/icon_detect/model.pt")
|
||
self.yolo.to(self.device)
|
||
|
||
# Florence-2 图标描述模型
|
||
self.processor = None
|
||
self.model = None
|
||
if enable_caption:
|
||
self.processor = AutoProcessor.from_pretrained(
|
||
"microsoft/Florence-2-base", trust_remote_code=True
|
||
)
|
||
self.model = AutoModelForCausalLM.from_pretrained(
|
||
f"{model_dir}/icon_caption",
|
||
torch_dtype=torch.float16 if self.device.type != "cpu" else torch.float32,
|
||
trust_remote_code=True,
|
||
).to(self.device)
|
||
|
||
# EasyOCR 文字识别模型
|
||
self.ocr = None
|
||
if enable_ocr:
|
||
langs = ocr_languages or ["ch_sim", "en"]
|
||
model_storage_directory = os.path.join(model_dir, "easyocr_models")
|
||
print(f"📄 EasyOCR 模型目录: {model_storage_directory}")
|
||
|
||
# EasyOCR recognition.py 内部 DataLoader 将 pin_memory 写死为 True
|
||
# 但 MPS 设备不支持 pin_memory，通过 monkey-patch 消除警告
|
||
if self.device.type == "mps":
|
||
try:
|
||
import easyocr.recognition as _easyocr_recognition
|
||
import torch.utils.data as _torch_data
|
||
_original_DataLoader = _torch_data.DataLoader
|
||
def _patched_DataLoader(*args, **kwargs):
|
||
kwargs['pin_memory'] = False
|
||
return _original_DataLoader(*args, **kwargs)
|
||
_easyocr_recognition.DataLoader = _patched_DataLoader
|
||
except Exception:
|
||
pass # 补丁失败不影响功能
|
||
|
||
self.ocr = easyocr.Reader(
|
||
langs,
|
||
gpu=self.device.type == "cuda",
|
||
model_storage_directory=model_storage_directory,
|
||
download_enabled=True
|
||
)
|
||
|
||
# box annotator
|
||
self.annotator = BoxAnnotator()
|
||
|
||
def __call__(self, data: Dict[str, Any]) -> Any:
|
||
data = data.pop("inputs")
|
||
|
||
# read image
|
||
image = load_image(data["image"])
|
||
|
||
# 1. OCR Step（如果启用）
|
||
ocr_texts, ocr_bboxes = [], []
|
||
if self.enable_ocr and self.ocr is not None:
|
||
ocr_texts, ocr_bboxes = self.check_ocr_bboxes(
|
||
image,
|
||
out_format="xyxy",
|
||
ocr_kwargs={"text_threshold": 0.4},
|
||
)
|
||
|
||
# 2. YOLO + Caption Step
|
||
annotated_image, filtered_bboxes_out = self.get_som_labeled_img(
|
||
image,
|
||
image_size=data.get("image_size", None),
|
||
ocr_texts=ocr_texts,
|
||
ocr_bboxes=ocr_bboxes,
|
||
bbox_threshold=data.get("bbox_threshold", 0.05),
|
||
iou_threshold=data.get("iou_threshold", 0.5),
|
||
)
|
||
return {
|
||
"image": annotated_image,
|
||
"bboxes": filtered_bboxes_out,
|
||
}
|
||
|
||
def check_ocr_bboxes(
|
||
self,
|
||
image: ImageType,
|
||
out_format: Literal["xywh", "xyxy"] = "xywh",
|
||
ocr_kwargs: Optional[Dict[str, Any]] = {},
|
||
) -> Tuple[List[str], List[List[int]]]:
|
||
# 🔧 修复点 1：RBGA -> RGBA
|
||
if image.mode == "RGBA":
|
||
image = image.convert("RGB")
|
||
|
||
result = self.ocr.readtext(np.array(image), **ocr_kwargs) # type: ignore
|
||
texts = [str(item[1]) for item in result]
|
||
bboxes = [
|
||
self.coordinates_to_bbox(item[0], format=out_format) for item in result
|
||
]
|
||
return (texts, bboxes)
|
||
|
||
@staticmethod
|
||
def coordinates_to_bbox(
|
||
coordinates: np.ndarray, format: Literal["xywh", "xyxy"] = "xywh"
|
||
) -> List[int]:
|
||
if format == "xywh":
|
||
return [
|
||
int(coordinates[0][0]),
|
||
int(coordinates[0][1]),
|
||
int(coordinates[2][0] - coordinates[0][0]),
|
||
int(coordinates[2][1] - coordinates[0][1]),
|
||
]
|
||
elif format == "xyxy":
|
||
return [
|
||
int(coordinates[0][0]),
|
||
int(coordinates[0][1]),
|
||
int(coordinates[2][0]),
|
||
int(coordinates[2][1]),
|
||
]
|
||
else:
|
||
raise ValueError(f"Unsupported format: {format}")
|
||
|
||
@staticmethod
|
||
def bbox_area(bbox: List[int], w: int, h: int) -> int:
|
||
bbox = [bbox[0] * w, bbox[1] * h, bbox[2] * w, bbox[3] * h]
|
||
return (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
|
||
|
||
@staticmethod
|
||
def remove_bbox_overlap(
|
||
xyxy_bboxes: List[Dict[str, Any]],
|
||
ocr_bboxes: Optional[List[Dict[str, Any]]] = None,
|
||
iou_threshold: Optional[float] = 0.7,
|
||
) -> List[Dict[str, Any]]:
|
||
filtered_bboxes = []
|
||
if ocr_bboxes is not None:
|
||
filtered_bboxes.extend(ocr_bboxes)
|
||
|
||
for i, bbox_outter in enumerate(xyxy_bboxes):
|
||
bbox_left = bbox_outter["bbox"]
|
||
valid_bbox = True
|
||
|
||
for j, bbox_inner in enumerate(xyxy_bboxes):
|
||
if i == j:
|
||
continue
|
||
|
||
bbox_right = bbox_inner["bbox"]
|
||
if (
|
||
intersection_over_union(
|
||
bbox_left,
|
||
bbox_right,
|
||
)
|
||
> iou_threshold # type: ignore
|
||
) and (area(bbox_left) > area(bbox_right)):
|
||
valid_bbox = False
|
||
break
|
||
|
||
if valid_bbox is False:
|
||
continue
|
||
|
||
if ocr_bboxes is None:
|
||
filtered_bboxes.append(bbox_outter)
|
||
continue
|
||
|
||
box_added = False
|
||
ocr_labels = []
|
||
for ocr_bbox in ocr_bboxes:
|
||
if not box_added:
|
||
bbox_right = ocr_bbox["bbox"]
|
||
# 检查是否重叠（合并 Icon 和 OCR 文字）
|
||
if overlap(bbox_right, bbox_left):
|
||
try:
|
||
ocr_labels.append(ocr_bbox["content"])
|
||
if ocr_bbox in filtered_bboxes:
|
||
filtered_bboxes.remove(ocr_bbox)
|
||
except Exception as e:
|
||
self.logger.error(f"Error removing bbox overlap: {e}")
|
||
continue
|
||
elif overlap(bbox_left, bbox_right):
|
||
box_added = True
|
||
break
|
||
|
||
if not box_added:
|
||
filtered_bboxes.append(
|
||
{
|
||
"type": "icon",
|
||
"bbox": bbox_outter["bbox"],
|
||
"interactivity": True,
|
||
"content": " ".join(ocr_labels) if ocr_labels else None,
|
||
}
|
||
)
|
||
|
||
return filtered_bboxes
|
||
|
||
def get_som_labeled_img(
|
||
self,
|
||
image: ImageType,
|
||
image_size: Optional[Dict[Literal["w", "h"], int]] = None,
|
||
ocr_texts: Optional[List[str]] = None,
|
||
ocr_bboxes: Optional[List[List[int]]] = None,
|
||
bbox_threshold: float = 0.05,
|
||
iou_threshold: Optional[float] = None,
|
||
caption_prompt: Optional[str] = None,
|
||
caption_batch_size: int = 64,
|
||
) -> Tuple[str, List[Dict[str, Any]]]:
|
||
if image.mode == "RGBA":
|
||
image = image.convert("RGB")
|
||
|
||
w, h = image.size
|
||
image_np = np.asarray(image)
|
||
|
||
# YOLO 检测（如果启用）
|
||
xyxy_bboxes_raw = torch.tensor([]) # 默认空
|
||
if self.enable_yolo and self.yolo is not None:
|
||
if image_size is None:
|
||
imgsz = [h, w]
|
||
else:
|
||
imgsz = [image_size.get("h", h), image_size.get("w", w)]
|
||
|
||
out = self.yolo.predict(
|
||
image,
|
||
imgsz=imgsz,
|
||
conf=bbox_threshold,
|
||
iou=iou_threshold or 0.7,
|
||
verbose=False,
|
||
)[0]
|
||
|
||
if out.boxes is not None:
|
||
xyxy_bboxes_raw = out.boxes.xyxy
|
||
xyxy_bboxes_raw = xyxy_bboxes_raw / torch.Tensor([w, h, w, h]).to(xyxy_bboxes_raw.device)
|
||
|
||
# 处理 OCR 检测框
|
||
ocr_bboxes_normalized = []
|
||
if ocr_bboxes:
|
||
ocr_bboxes_tensor = torch.tensor(ocr_bboxes) / torch.Tensor([w, h, w, h])
|
||
ocr_bboxes_normalized = ocr_bboxes_tensor.tolist()
|
||
|
||
ocr_bbox_dicts = [
|
||
{
|
||
"type": "text",
|
||
"bbox": bbox,
|
||
"interactivity": False,
|
||
"content": text,
|
||
"source": "box_ocr_content_ocr",
|
||
}
|
||
for bbox, text in zip(ocr_bboxes_normalized, ocr_texts or [])
|
||
if self.bbox_area(bbox, w, h) > 0
|
||
]
|
||
|
||
# 处理 YOLO 检测框
|
||
yolo_bbox_dicts = [
|
||
{
|
||
"type": "icon",
|
||
"bbox": bbox,
|
||
"interactivity": True,
|
||
"content": None,
|
||
"source": "box_yolo_content_yolo",
|
||
}
|
||
for bbox in xyxy_bboxes_raw.tolist()
|
||
if self.bbox_area(bbox, w, h) > 0
|
||
]
|
||
|
||
filtered_bboxes = self.remove_bbox_overlap(
|
||
xyxy_bboxes=yolo_bbox_dicts,
|
||
ocr_bboxes=ocr_bbox_dicts if ocr_bbox_dicts else None,
|
||
iou_threshold=iou_threshold or 0.7,
|
||
)
|
||
|
||
filtered_bboxes_out = sorted(
|
||
filtered_bboxes, key=lambda x: x["content"] is None
|
||
)
|
||
starting_idx = next(
|
||
(
|
||
idx
|
||
for idx, bbox in enumerate(filtered_bboxes_out)
|
||
if bbox["content"] is None
|
||
),
|
||
-1,
|
||
)
|
||
|
||
filtered_bboxes = torch.tensor([box["bbox"] for box in filtered_bboxes_out])
|
||
|
||
# 如果启用了 Caption 且有需要描述的图标
|
||
if starting_idx != -1 and self.enable_caption and self.model is not None:
|
||
non_ocr_bboxes = filtered_bboxes[starting_idx:]
|
||
bbox_images = []
|
||
for _, coordinates in enumerate(non_ocr_bboxes):
|
||
try:
|
||
xmin, xmax = (
|
||
int(coordinates[0] * image_np.shape[1]),
|
||
int(coordinates[2] * image_np.shape[1]),
|
||
)
|
||
ymin, ymax = (
|
||
int(coordinates[1] * image_np.shape[0]),
|
||
int(coordinates[3] * image_np.shape[0]),
|
||
)
|
||
cropped_image = image_np[ymin:ymax, xmin:xmax, :]
|
||
cropped_image = cv2.resize(cropped_image, (64, 64))
|
||
bbox_images.append(ToPILImage()(cropped_image))
|
||
except Exception as e:
|
||
self.logger.error(f"Error cropping bbox: {e}")
|
||
continue
|
||
|
||
if caption_prompt is None:
|
||
caption_prompt = "<CAPTION>"
|
||
|
||
captions = []
|
||
for idx in range(0, len(bbox_images), caption_batch_size): # type: ignore
|
||
batch = bbox_images[idx : idx + caption_batch_size] # type: ignore
|
||
if not batch: break
|
||
inputs = self.processor(
|
||
images=batch,
|
||
text=[caption_prompt] * len(batch),
|
||
return_tensors="pt",
|
||
do_resize=False,
|
||
)
|
||
if self.device.type in {"cuda", "mps"}:
|
||
inputs = inputs.to(device=self.device, dtype=torch.float16)
|
||
|
||
with torch.inference_mode():
|
||
generated_ids = self.model.generate(
|
||
input_ids=inputs["input_ids"],
|
||
pixel_values=inputs["pixel_values"],
|
||
max_new_tokens=20,
|
||
num_beams=1,
|
||
do_sample=False,
|
||
early_stopping=False,
|
||
)
|
||
|
||
generated_texts = self.processor.batch_decode(
|
||
generated_ids, skip_special_tokens=True
|
||
)
|
||
captions.extend([text.strip() for text in generated_texts])
|
||
|
||
ocr_texts = [f"Text Box ID {idx}: {text}" for idx, text in enumerate(ocr_texts)] # type: ignore
|
||
for _, bbox in enumerate(filtered_bboxes_out):
|
||
if bbox["content"] is None and captions:
|
||
bbox["content"] = captions.pop(0)
|
||
|
||
filtered_bboxes = box_convert(
|
||
boxes=filtered_bboxes, in_fmt="xyxy", out_fmt="cxcywh"
|
||
)
|
||
|
||
annotated_image = image_np.copy()
|
||
bboxes_annotate = filtered_bboxes * torch.Tensor([w, h, w, h])
|
||
xyxy_annotate = box_convert(
|
||
bboxes_annotate, in_fmt="cxcywh", out_fmt="xyxy"
|
||
).numpy()
|
||
detections = Detections(xyxy=xyxy_annotate)
|
||
labels = [str(idx) for idx in range(bboxes_annotate.shape[0])]
|
||
|
||
annotated_image = self.annotator.annotate(
|
||
scene=annotated_image,
|
||
detections=detections,
|
||
labels=labels,
|
||
image_size=(w, h),
|
||
)
|
||
assert w == annotated_image.shape[1] and h == annotated_image.shape[0]
|
||
|
||
out_image = Image.fromarray(annotated_image)
|
||
out_buffer = io.BytesIO()
|
||
out_image.save(out_buffer, format="PNG")
|
||
encoded_image = base64.b64encode(out_buffer.getvalue()).decode("ascii")
|
||
|
||
return encoded_image, filtered_bboxes_out
|
||
|
||
|
||
def area(bbox: List[int]) -> int:
|
||
return (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
|
||
|
||
|
||
def intersection_area(bbox_left: List[int], bbox_right: List[int]) -> int:
|
||
# 计算两个 bbox 的交集面积
|
||
# 交集的左边界取两者的最大值，右边界取两者的最小值
|
||
x_overlap = max(0, min(bbox_left[2], bbox_right[2]) - max(bbox_left[0], bbox_right[0]))
|
||
y_overlap = max(0, min(bbox_left[3], bbox_right[3]) - max(bbox_left[1], bbox_right[1]))
|
||
return x_overlap * y_overlap
|
||
|
||
|
||
def intersection_over_union(bbox_left: List[int], bbox_right: List[int]) -> float:
|
||
intersection = intersection_area(bbox_left, bbox_right)
|
||
bbox_left_area = area(bbox_left)
|
||
bbox_right_area = area(bbox_right)
|
||
union = bbox_left_area + bbox_right_area - intersection + 1e-6
|
||
|
||
ratio_left, ratio_right = 0, 0
|
||
if bbox_left_area > 0 and bbox_right_area > 0:
|
||
ratio_left = intersection / bbox_left_area
|
||
ratio_right = intersection / bbox_right_area
|
||
return max(intersection / union, ratio_left, ratio_right)
|
||
|
||
|
||
def overlap(bbox_left: List[int], bbox_right: List[int]) -> bool:
|
||
intersection = intersection_area(bbox_left, bbox_right)
|
||
ratio_left = intersection / area(bbox_left)
|
||
# 🔧 修复点 3：从 0.80 降低到 0.50
|
||
# 作用：只要 OCR 文字框有一半在 YOLO 图标框内，就认为它们是一体的，
|
||
# 避免因为框对齐不准导致 Icon 框被丢弃或产生两个框
|
||
return ratio_left > 0.50
|
||
|
||
|
||
class BoxAnnotator:
|
||
def __init__(
|
||
self,
|
||
color: Union[Color, ColorPalette] = ColorPalette.DEFAULT, # type: ignore
|
||
thickness: int = 3,
|
||
text_color: Color = Color.BLACK, # type: ignore
|
||
text_scale: float = 0.5,
|
||
text_thickness: int = 2,
|
||
text_padding: int = 10,
|
||
avoid_overlap: bool = True,
|
||
):
|
||
self.color: Union[Color, ColorPalette] = color
|
||
self.thickness: int = thickness
|
||
self.text_color: Color = text_color
|
||
self.text_scale: float = text_scale
|
||
self.text_thickness: int = text_thickness
|
||
self.text_padding: int = text_padding
|
||
self.avoid_overlap: bool = avoid_overlap
|
||
|
||
def annotate(
|
||
self,
|
||
scene: np.ndarray,
|
||
detections: Detections,
|
||
labels: Optional[List[str]] = None,
|
||
skip_label: bool = False,
|
||
image_size: Optional[Tuple[int, int]] = None,
|
||
) -> np.ndarray:
|
||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||
for i in range(len(detections)):
|
||
x1, y1, x2, y2 = detections.xyxy[i].astype(int)
|
||
class_id = (
|
||
detections.class_id[i] if detections.class_id is not None else None
|
||
)
|
||
idx = class_id if class_id is not None else i
|
||
color = (
|
||
self.color.by_idx(idx)
|
||
if isinstance(self.color, ColorPalette)
|
||
else self.color
|
||
)
|
||
cv2.rectangle(
|
||
img=scene,
|
||
pt1=(x1, y1),
|
||
pt2=(x2, y2),
|
||
color=color.as_bgr(),
|
||
thickness=self.thickness,
|
||
)
|
||
if skip_label:
|
||
continue
|
||
|
||
text = (
|
||
f"{class_id}"
|
||
if (labels is None or len(detections) != len(labels))
|
||
else labels[i]
|
||
)
|
||
|
||
text_width, text_height = cv2.getTextSize(
|
||
text=text,
|
||
fontFace=font,
|
||
fontScale=self.text_scale,
|
||
thickness=self.text_thickness,
|
||
)[0]
|
||
|
||
if not self.avoid_overlap:
|
||
text_x = x1 + self.text_padding
|
||
text_y = y1 - self.text_padding
|
||
|
||
text_background_x1 = x1
|
||
text_background_y1 = y1 - 2 * self.text_padding - text_height
|
||
|
||
text_background_x2 = x1 + 2 * self.text_padding + text_width
|
||
text_background_y2 = y1
|
||
else:
|
||
(
|
||
text_x,
|
||
text_y,
|
||
text_background_x1,
|
||
text_background_y1,
|
||
text_background_x2,
|
||
text_background_y2,
|
||
) = self.get_optimal_label_pos(
|
||
self.text_padding,
|
||
text_width,
|
||
text_height,
|
||
x1,
|
||
y1,
|
||
x2,
|
||
y2,
|
||
detections,
|
||
image_size,
|
||
)
|
||
|
||
cv2.rectangle(
|
||
img=scene,
|
||
pt1=(text_background_x1, text_background_y1),
|
||
pt2=(text_background_x2, text_background_y2),
|
||
color=color.as_bgr(),
|
||
thickness=cv2.FILLED,
|
||
)
|
||
box_color = color.as_rgb()
|
||
luminance = (
|
||
0.299 * box_color[0] + 0.587 * box_color[1] + 0.114 * box_color[2]
|
||
)
|
||
text_color = (0, 0, 0) if luminance > 160 else (255, 255, 255)
|
||
cv2.putText(
|
||
img=scene,
|
||
text=text,
|
||
org=(text_x, text_y),
|
||
fontFace=font,
|
||
fontScale=self.text_scale,
|
||
color=text_color,
|
||
thickness=self.text_thickness,
|
||
lineType=cv2.LINE_AA,
|
||
)
|
||
return scene
|
||
|
||
@staticmethod
|
||
def get_optimal_label_pos(
|
||
text_padding, text_width, text_height, x1, y1, x2, y2, detections, image_size
|
||
):
|
||
def get_is_overlap(
|
||
detections,
|
||
text_background_x1,
|
||
text_background_y1,
|
||
text_background_x2,
|
||
text_background_y2,
|
||
image_size,
|
||
):
|
||
is_overlap = False
|
||
for i in range(len(detections)):
|
||
detection = detections.xyxy[i].astype(int)
|
||
if (
|
||
intersection_over_union(
|
||
[
|
||
text_background_x1,
|
||
text_background_y1,
|
||
text_background_x2,
|
||
text_background_y2,
|
||
],
|
||
detection,
|
||
)
|
||
> 0.3
|
||
):
|
||
is_overlap = True
|
||
break
|
||
if (
|
||
text_background_x1 < 0
|
||
or text_background_x2 > image_size[0]
|
||
or text_background_y1 < 0
|
||
or text_background_y2 > image_size[1]
|
||
):
|
||
is_overlap = True
|
||
return is_overlap
|
||
|
||
text_x = x1 + text_padding
|
||
text_y = y1 - text_padding
|
||
|
||
text_background_x1 = x1
|
||
text_background_y1 = y1 - 2 * text_padding - text_height
|
||
|
||
text_background_x2 = x1 + 2 * text_padding + text_width
|
||
text_background_y2 = y1
|
||
is_overlap = get_is_overlap(
|
||
detections,
|
||
text_background_x1,
|
||
text_background_y1,
|
||
text_background_x2,
|
||
text_background_y2,
|
||
image_size,
|
||
)
|
||
if not is_overlap:
|
||
return (
|
||
text_x,
|
||
text_y,
|
||
text_background_x1,
|
||
text_background_y1,
|
||
text_background_x2,
|
||
text_background_y2,
|
||
)
|
||
|
||
text_x = x1 - text_padding - text_width
|
||
text_y = y1 + text_padding + text_height
|
||
|
||
text_background_x1 = x1 - 2 * text_padding - text_width
|
||
text_background_y1 = y1
|
||
|
||
text_background_x2 = x1
|
||
text_background_y2 = y1 + 2 * text_padding + text_height
|
||
is_overlap = get_is_overlap(
|
||
detections,
|
||
text_background_x1,
|
||
text_background_y1,
|
||
text_background_x2,
|
||
text_background_y2,
|
||
image_size,
|
||
)
|
||
if not is_overlap:
|
||
return (
|
||
text_x,
|
||
text_y,
|
||
text_background_x1,
|
||
text_background_y1,
|
||
text_background_x2,
|
||
text_background_y2,
|
||
)
|
||
|
||
text_x = x2 + text_padding
|
||
text_y = y1 + text_padding + text_height
|
||
|
||
text_background_x1 = x2
|
||
text_background_y1 = y1
|
||
|
||
text_background_x2 = x2 + 2 * text_padding + text_width
|
||
text_background_y2 = y1 + 2 * text_padding + text_height
|
||
|
||
is_overlap = get_is_overlap(
|
||
detections,
|
||
text_background_x1,
|
||
text_background_y1,
|
||
text_background_x2,
|
||
text_background_y2,
|
||
image_size,
|
||
)
|
||
if not is_overlap:
|
||
return (
|
||
text_x,
|
||
text_y,
|
||
text_background_x1,
|
||
text_background_y1,
|
||
text_background_x2,
|
||
text_background_y2,
|
||
)
|
||
|
||
text_x = x2 - text_padding - text_width
|
||
text_y = y1 - text_padding
|
||
|
||
text_background_x1 = x2 - 2 * text_padding - text_width
|
||
text_background_y1 = y1 - 2 * text_padding - text_height
|
||
|
||
text_background_x2 = x2
|
||
text_background_y2 = y1
|
||
|
||
is_overlap = get_is_overlap(
|
||
detections,
|
||
text_background_x1,
|
||
text_background_y1,
|
||
text_background_x2,
|
||
text_background_y2,
|
||
image_size,
|
||
)
|
||
if not is_overlap:
|
||
return (
|
||
text_x,
|
||
text_y,
|
||
text_background_x1,
|
||
text_background_y1,
|
||
text_background_x2,
|
||
text_background_y2,
|
||
)
|
||
|
||
return (
|
||
text_x,
|
||
text_y,
|
||
text_background_x1,
|
||
text_background_y1,
|
||
text_background_x2,
|
||
text_background_y2,
|
||
) |