452 lines
13 KiB
Python
452 lines
13 KiB
Python
from pathlib import Path
|
|
from typing import List, Dict
|
|
import base64
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import time
|
|
import xml.etree.ElementTree as ET
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
from app.config import settings
|
|
|
|
|
|
def read_classes(path: str, keep_empty: bool = False) -> List[str]:
|
|
"""
|
|
支持两种类别文件格式：
|
|
1. 每行一个类别；
|
|
2. 用英文分号 ; 分隔类别。
|
|
|
|
keep_empty=True 时，会保留空类别。
|
|
你的分类文件里有 ;;，表示该类别不显示，所以分类类别必须保留空项。
|
|
"""
|
|
p = Path(path)
|
|
if not p.exists():
|
|
return []
|
|
|
|
text = p.read_text(encoding="utf-8").strip()
|
|
|
|
if ";" in text:
|
|
parts = text.split(";")
|
|
if keep_empty:
|
|
return [x.strip() for x in parts]
|
|
return [x.strip() for x in parts if x.strip()]
|
|
|
|
lines = text.splitlines()
|
|
if keep_empty:
|
|
return [x.strip() for x in lines]
|
|
return [x.strip() for x in lines if x.strip()]
|
|
|
|
|
|
def decode_image_from_bytes(data: bytes) -> np.ndarray:
|
|
arr = np.frombuffer(data, dtype=np.uint8)
|
|
img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
|
|
if img is None:
|
|
raise ValueError("图片解码失败，请确认上传的是有效图片文件")
|
|
return img
|
|
|
|
|
|
def encode_image_to_base64(img_bgr: np.ndarray) -> str:
|
|
ok, buf = cv2.imencode(".jpg", img_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), 90])
|
|
if not ok:
|
|
raise ValueError("图片编码失败")
|
|
return base64.b64encode(buf.tobytes()).decode("utf-8")
|
|
|
|
|
|
def clip_box(box, width: int, height: int):
|
|
x1, y1, x2, y2 = box
|
|
x1 = max(0, min(float(x1), width - 1))
|
|
y1 = max(0, min(float(y1), height - 1))
|
|
x2 = max(0, min(float(x2), width - 1))
|
|
y2 = max(0, min(float(y2), height - 1))
|
|
return [x1, y1, x2, y2]
|
|
|
|
|
|
def crop_with_padding(img_bgr: np.ndarray, box, padding_ratio: float = 0.08) -> np.ndarray:
|
|
h, w = img_bgr.shape[:2]
|
|
x1, y1, x2, y2 = clip_box(box, w, h)
|
|
|
|
bw = x2 - x1
|
|
bh = y2 - y1
|
|
pad_x = bw * padding_ratio
|
|
pad_y = bh * padding_ratio
|
|
|
|
nx1 = int(max(0, x1 - pad_x))
|
|
ny1 = int(max(0, y1 - pad_y))
|
|
nx2 = int(min(w, x2 + pad_x))
|
|
ny2 = int(min(h, y2 + pad_y))
|
|
|
|
if nx2 <= nx1 or ny2 <= ny1:
|
|
return img_bgr.copy()
|
|
|
|
return img_bgr[ny1:ny2, nx1:nx2].copy()
|
|
|
|
|
|
def stable_color_for_class(class_id: int, class_name: str = "") -> tuple[int, int, int]:
|
|
"""
|
|
根据类别 ID 和类别名生成稳定颜色。
|
|
同一物种每次运行颜色一致。
|
|
OpenCV 使用 BGR 顺序。
|
|
"""
|
|
key = f"{class_id}_{class_name}".encode("utf-8")
|
|
digest = hashlib.md5(key).hexdigest()
|
|
|
|
b = int(digest[0:2], 16)
|
|
g = int(digest[2:4], 16)
|
|
r = int(digest[4:6], 16)
|
|
|
|
# 避免颜色太暗
|
|
b = 80 + b % 176
|
|
g = 80 + g % 176
|
|
r = 80 + r % 176
|
|
|
|
return b, g, r
|
|
|
|
|
|
def bgr_to_rgb(color_bgr: tuple[int, int, int]) -> tuple[int, int, int]:
|
|
b, g, r = color_bgr
|
|
return r, g, b
|
|
|
|
|
|
def get_display_label(item: Dict) -> str:
|
|
"""
|
|
可视化时只显示物种中文名。
|
|
如果分类文件中该类别是空字符串，则不显示文字。
|
|
"""
|
|
name = str(item.get("cls_class_name", "") or "").strip()
|
|
return name
|
|
|
|
|
|
def find_chinese_font() -> str | None:
|
|
"""
|
|
自动寻找中文字体。
|
|
可优先通过环境变量 FONT_PATH 指定。
|
|
"""
|
|
env_font = os.getenv("FONT_PATH", "").strip()
|
|
if env_font and Path(env_font).exists():
|
|
return env_font
|
|
|
|
candidates = [
|
|
"/data2/zhangw_data/Bird_detect_cascade_service/fonts/NotoSansCJK-Regular.ttc",
|
|
"/data2/zhangw_data/Bird_detect_cascade_service/fonts/NotoSansCJK-Regular.otf",
|
|
"/data2/zhangw_data/Bird_detect_cascade_service/fonts/SourceHanSansSC-Regular.otf",
|
|
"/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc",
|
|
"/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.otf",
|
|
"/usr/share/fonts/truetype/noto/NotoSansCJK-Regular.ttc",
|
|
"/usr/share/fonts/truetype/noto/NotoSansCJK-Regular.otf",
|
|
"/usr/share/fonts/truetype/wqy/wqy-microhei.ttc",
|
|
"/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc",
|
|
"/usr/share/fonts/truetype/arphic/uming.ttc",
|
|
"/usr/share/fonts/truetype/arphic/ukai.ttc",
|
|
"/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
|
|
]
|
|
|
|
for p in candidates:
|
|
if Path(p).exists():
|
|
return p
|
|
|
|
return None
|
|
|
|
|
|
_FONT_CACHE: dict[int, ImageFont.FreeTypeFont | ImageFont.ImageFont] = {}
|
|
|
|
|
|
def get_font(font_size: int = 26):
|
|
"""
|
|
获取支持中文的字体。
|
|
如果找不到中文字体，会退回默认字体，但默认字体仍可能无法显示中文。
|
|
"""
|
|
if font_size in _FONT_CACHE:
|
|
return _FONT_CACHE[font_size]
|
|
|
|
font_path = find_chinese_font()
|
|
|
|
if font_path:
|
|
font = ImageFont.truetype(font_path, font_size)
|
|
else:
|
|
font = ImageFont.load_default()
|
|
|
|
_FONT_CACHE[font_size] = font
|
|
return font
|
|
|
|
|
|
def draw_chinese_text_with_background(
|
|
img_bgr: np.ndarray,
|
|
text: str,
|
|
left: int,
|
|
top: int,
|
|
color_bgr: tuple[int, int, int],
|
|
font_size: int = 26,
|
|
) -> np.ndarray:
|
|
"""
|
|
使用 PIL 绘制中文标签。
|
|
输入输出都是 OpenCV BGR 图像。
|
|
|
|
这个版本解决中文文字在色块中偏下的问题：
|
|
1. 先计算文字真实边界；
|
|
2. 再根据背景框中心点反推文字绘制位置；
|
|
3. 保证中文视觉上垂直居中。
|
|
"""
|
|
if not text:
|
|
return img_bgr
|
|
|
|
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
|
pil_img = Image.fromarray(img_rgb)
|
|
draw = ImageDraw.Draw(pil_img)
|
|
|
|
font = get_font(font_size)
|
|
color_rgb = bgr_to_rgb(color_bgr)
|
|
|
|
# 计算文字真实边界。注意：中文字体 bbox 的 y 方向通常不是从 0 开始。
|
|
bbox = draw.textbbox((0, 0), text, font=font)
|
|
text_left = bbox[0]
|
|
text_top = bbox[1]
|
|
text_right = bbox[2]
|
|
text_bottom = bbox[3]
|
|
|
|
text_w = text_right - text_left
|
|
text_h = text_bottom - text_top
|
|
|
|
pad_x = 6
|
|
pad_y = 4
|
|
|
|
img_w, img_h = pil_img.size
|
|
|
|
x1 = max(0, min(int(left), img_w - 1))
|
|
y1 = max(0, min(int(top), img_h - 1))
|
|
|
|
bg_w = text_w + 2 * pad_x
|
|
bg_h = text_h + 2 * pad_y
|
|
|
|
# 默认把标签放在检测框上方；如果上方空间不够，就放在框内部顶端。
|
|
bg_left = x1
|
|
if y1 - bg_h >= 0:
|
|
bg_top = y1 - bg_h
|
|
else:
|
|
bg_top = y1
|
|
|
|
bg_right = min(img_w - 1, bg_left + bg_w)
|
|
bg_bottom = min(img_h - 1, bg_top + bg_h)
|
|
|
|
# 画背景色块
|
|
draw.rectangle(
|
|
[bg_left, bg_top, bg_right, bg_bottom],
|
|
fill=color_rgb,
|
|
)
|
|
|
|
# 关键修正：
|
|
# 让文字真实 bbox 的中心对齐背景框中心，而不是简单 bg_top + pad_y。
|
|
bg_center_x = (bg_left + bg_right) / 2
|
|
bg_center_y = (bg_top + bg_bottom) / 2
|
|
|
|
text_draw_x = bg_center_x - text_w / 2 - text_left
|
|
text_draw_y = bg_center_y - text_h / 2 - text_top
|
|
|
|
draw.text(
|
|
(text_draw_x, text_draw_y),
|
|
text,
|
|
font=font,
|
|
fill=(0, 0, 0),
|
|
)
|
|
|
|
out_rgb = np.array(pil_img)
|
|
out_bgr = cv2.cvtColor(out_rgb, cv2.COLOR_RGB2BGR)
|
|
return out_bgr
|
|
|
|
def draw_results(img_bgr: np.ndarray, results: List[dict]) -> np.ndarray:
|
|
"""
|
|
只画类别名，不画置信度。
|
|
中文文字使用 PIL 绘制，避免 OpenCV 中文乱码。
|
|
空类别名只画框，不显示文字。
|
|
"""
|
|
out = img_bgr.copy()
|
|
|
|
for item in results:
|
|
x1, y1, x2, y2 = [int(round(v)) for v in item["bbox_xyxy"]]
|
|
|
|
cls_id = int(item.get("cls_class_id", -1))
|
|
cls_name = str(item.get("cls_class_name", "") or "")
|
|
label = get_display_label(item)
|
|
|
|
if settings.uniform_species_color:
|
|
color = (0, 190, 0)
|
|
else:
|
|
color = stable_color_for_class(cls_id, cls_name)
|
|
|
|
# 先用 OpenCV 画框
|
|
cv2.rectangle(out, (x1, y1), (x2, y2), color, 2)
|
|
|
|
# 再用 PIL 画中文
|
|
if label:
|
|
# 根据图像大小自适应字体大小，避免无人机大图上文字太小
|
|
h, w = out.shape[:2]
|
|
font_size = max(22, int(min(w, h) * 0.018))
|
|
font_size = min(font_size, 42)
|
|
|
|
out = draw_chinese_text_with_background(
|
|
img_bgr=out,
|
|
text=label,
|
|
left=x1,
|
|
top=y1,
|
|
color_bgr=color,
|
|
font_size=font_size,
|
|
)
|
|
|
|
return out
|
|
|
|
|
|
def now_timestamp_minute() -> str:
|
|
return time.strftime("%Y%m%d%H%M")
|
|
|
|
|
|
def safe_stem(filename: str) -> str:
|
|
name = Path(filename).stem
|
|
name = name.replace(" ", "_")
|
|
return name
|
|
|
|
|
|
def make_unique_base(output_dir: Path, timestamp: str, original_filename: str) -> str:
|
|
"""
|
|
生成 时间戳_原文件名。
|
|
如果同一分钟内重复处理同一个文件，则自动加 _01、_02。
|
|
"""
|
|
stem = safe_stem(original_filename)
|
|
base = f"{timestamp}_{stem}"
|
|
|
|
candidate = base
|
|
idx = 1
|
|
while (
|
|
(output_dir / f"{candidate}.json").exists()
|
|
or (output_dir / f"{candidate}.xml").exists()
|
|
or (output_dir / f"{candidate}.jpg").exists()
|
|
or (output_dir / f"{candidate}.mp4").exists()
|
|
):
|
|
candidate = f"{base}_{idx:02d}"
|
|
idx += 1
|
|
|
|
return candidate
|
|
|
|
|
|
def save_json(path: Path, data: Dict):
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
|
|
|
|
def indent_xml(elem, level: int = 0):
|
|
i = "\n" + level * " "
|
|
if len(elem):
|
|
if not elem.text or not elem.text.strip():
|
|
elem.text = i + " "
|
|
for child in elem:
|
|
indent_xml(child, level + 1)
|
|
if not child.tail or not child.tail.strip():
|
|
child.tail = i
|
|
if level and (not elem.tail or not elem.tail.strip()):
|
|
elem.tail = i
|
|
|
|
|
|
def write_image_xml(
|
|
xml_path: Path,
|
|
image_filename: str,
|
|
image_path: str,
|
|
width: int,
|
|
height: int,
|
|
results: List[Dict],
|
|
):
|
|
"""
|
|
图片 XML，接近 Pascal VOC 格式。
|
|
空类别名为了避免 XML name 为空，会写成 class_<id>。
|
|
"""
|
|
root = ET.Element("annotation")
|
|
|
|
ET.SubElement(root, "filename").text = image_filename
|
|
ET.SubElement(root, "path").text = image_path
|
|
|
|
size = ET.SubElement(root, "size")
|
|
ET.SubElement(size, "width").text = str(width)
|
|
ET.SubElement(size, "height").text = str(height)
|
|
ET.SubElement(size, "depth").text = "3"
|
|
|
|
for item in results:
|
|
obj = ET.SubElement(root, "object")
|
|
|
|
cls_id = int(item.get("cls_class_id", -1))
|
|
cls_name = str(item.get("cls_class_name", "") or "").strip()
|
|
xml_name = cls_name if cls_name else f"class_{cls_id}"
|
|
|
|
ET.SubElement(obj, "name").text = xml_name
|
|
ET.SubElement(obj, "class_id").text = str(cls_id)
|
|
|
|
box = item["bbox_xyxy"]
|
|
bndbox = ET.SubElement(obj, "bndbox")
|
|
ET.SubElement(bndbox, "xmin").text = str(int(round(box[0])))
|
|
ET.SubElement(bndbox, "ymin").text = str(int(round(box[1])))
|
|
ET.SubElement(bndbox, "xmax").text = str(int(round(box[2])))
|
|
ET.SubElement(bndbox, "ymax").text = str(int(round(box[3])))
|
|
|
|
indent_xml(root)
|
|
tree = ET.ElementTree(root)
|
|
xml_path.parent.mkdir(parents=True, exist_ok=True)
|
|
tree.write(xml_path, encoding="utf-8", xml_declaration=True)
|
|
|
|
|
|
def write_video_xml(
|
|
xml_path: Path,
|
|
video_filename: str,
|
|
video_path: str,
|
|
width: int,
|
|
height: int,
|
|
fps: float,
|
|
total_frames: int,
|
|
sample_interval: int,
|
|
frame_results: List[Dict],
|
|
):
|
|
"""
|
|
视频 XML：记录每个处理帧上的目标框。
|
|
"""
|
|
root = ET.Element("video_annotation")
|
|
|
|
ET.SubElement(root, "filename").text = video_filename
|
|
ET.SubElement(root, "path").text = video_path
|
|
|
|
info = ET.SubElement(root, "video_info")
|
|
ET.SubElement(info, "width").text = str(width)
|
|
ET.SubElement(info, "height").text = str(height)
|
|
ET.SubElement(info, "fps").text = str(fps)
|
|
ET.SubElement(info, "total_frames").text = str(total_frames)
|
|
ET.SubElement(info, "sample_interval").text = str(sample_interval)
|
|
|
|
frames_node = ET.SubElement(root, "frames")
|
|
|
|
for frame_item in frame_results:
|
|
frame_node = ET.SubElement(frames_node, "frame")
|
|
ET.SubElement(frame_node, "frame_index").text = str(frame_item["frame_index"])
|
|
ET.SubElement(frame_node, "time_sec").text = str(frame_item["time_sec"])
|
|
|
|
objects_node = ET.SubElement(frame_node, "objects")
|
|
|
|
for item in frame_item.get("results", []):
|
|
obj = ET.SubElement(objects_node, "object")
|
|
|
|
cls_id = int(item.get("cls_class_id", -1))
|
|
cls_name = str(item.get("cls_class_name", "") or "").strip()
|
|
xml_name = cls_name if cls_name else f"class_{cls_id}"
|
|
|
|
ET.SubElement(obj, "name").text = xml_name
|
|
ET.SubElement(obj, "class_id").text = str(cls_id)
|
|
|
|
box = item["bbox_xyxy"]
|
|
bndbox = ET.SubElement(obj, "bndbox")
|
|
ET.SubElement(bndbox, "xmin").text = str(int(round(box[0])))
|
|
ET.SubElement(bndbox, "ymin").text = str(int(round(box[1])))
|
|
ET.SubElement(bndbox, "xmax").text = str(int(round(box[2])))
|
|
ET.SubElement(bndbox, "ymax").text = str(int(round(box[3])))
|
|
|
|
indent_xml(root)
|
|
tree = ET.ElementTree(root)
|
|
xml_path.parent.mkdir(parents=True, exist_ok=True)
|
|
tree.write(xml_path, encoding="utf-8", xml_declaration=True)
|