initial commit
This commit is contained in:
Ankle committed
commit
0bafbee31f
54 files changed
+6317
No files matched your search
+451
@@ -0,0 +1,451 @@
|
||||
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)
|
||||
Reference in new issue
Block a user