199 lines
8.6 KiB
Python
199 lines
8.6 KiB
Python
import cv2
|
||
import numpy as np
|
||
import os
|
||
from tqdm import tqdm
|
||
|
||
# -------------------------- 1. 核心配置（请根据需求修改） --------------------------
|
||
# 输入目录（图片和标签需一一对应，文件名相同，仅后缀不同）
|
||
INPUT_IMAGE_DIR = r"D:\DataPreHandler\data\train\images" # 原始图片目录
|
||
INPUT_LABEL_DIR = r"D:\DataPreHandler\data\train\labels" # 原始YOLO标签目录
|
||
# 输出目录（标注后的图片会保存在这里）
|
||
OUTPUT_IMAGE_DIR = r"D:\DataPreHandler\data\test\da\output2"
|
||
|
||
# 关键配置：类别ID与类别名称的映射（必须与你的YOLO训练类别顺序一致！）
|
||
CLASS_CONFIG = [
|
||
(0, "Abdomen", (0, 255, 0)),
|
||
(1, "Hips", (0, 255, 0)),
|
||
(2, "Chest", (0, 255, 0)),
|
||
(3, "vulva", (0, 255, 0)),
|
||
(4, "back", (0, 255, 0)),
|
||
(5, "penis", (0, 255, 0)),
|
||
(6, "Horror", (0, 255, 0))
|
||
]
|
||
|
||
# 绘制参数（可按需调整）
|
||
BOX_THICKNESS = 2 # 边界框线条厚度（像素）
|
||
FONT_FACE = cv2.FONT_HERSHEY_SIMPLEX # 字体类型
|
||
FONT_SCALE = 0.6 # 字体大小（根据图片尺寸调整）
|
||
FONT_THICKNESS = 1 # 字体线条厚度
|
||
TEXT_PADDING = 5 # 文字与边界框的间距（像素）
|
||
TEXT_BG_OPACITY = 0.7 # 文字背景的透明度（0-1，0为完全透明）
|
||
|
||
# 支持的图片格式（无需修改）
|
||
SUPPORTED_IMAGE_FORMATS = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff')
|
||
|
||
|
||
# -------------------------- 2. 工具函数 --------------------------
|
||
def yolo2pixel(yolo_coords, img_w, img_h):
|
||
"""
|
||
将YOLO相对坐标转换为图片像素坐标（边界框：x1, y1, x2, y2）
|
||
:param yolo_coords: YOLO坐标列表 [xc, yc, w, h]（相对值，0-1）
|
||
:param img_w: 图片宽度（像素）
|
||
:param img_h: 图片高度（像素）
|
||
:return: 像素坐标元组 (x1, y1, x2, y2)
|
||
"""
|
||
xc, yc, w, h = yolo_coords
|
||
# 计算边界框左上角和右下角坐标
|
||
x1 = int((xc - w / 2) * img_w)
|
||
y1 = int((yc - h / 2) * img_h)
|
||
x2 = int((xc + w / 2) * img_w)
|
||
y2 = int((yc + h / 2) * img_h)
|
||
# 确保坐标不超出图片范围
|
||
x1 = max(0, x1)
|
||
y1 = max(0, y1)
|
||
x2 = min(img_w, x2)
|
||
y2 = min(img_h, y2)
|
||
return x1, y1, x2, y2
|
||
|
||
|
||
def draw_annotation(img, bbox, class_name, color):
|
||
"""
|
||
在图片上绘制边界框和类别名称
|
||
:param img: 原始图片（OpenCV格式，BGR通道）
|
||
:param bbox: 像素坐标边界框 (x1, y1, x2, y2)
|
||
:param class_name: 类别名称（字符串）
|
||
:param color: 边界框和文字颜色（BGR元组，如 (0,255,0) 代表绿色）
|
||
:return: 标注后的图片
|
||
"""
|
||
img_h, img_w = img.shape[:2]
|
||
x1, y1, x2, y2 = bbox
|
||
|
||
# 1. 绘制边界框
|
||
cv2.rectangle(img, (x1, y1), (x2, y2), color, BOX_THICKNESS)
|
||
|
||
# 2. 计算文字尺寸（用于创建文字背景）
|
||
text_size, _ = cv2.getTextSize(class_name, FONT_FACE, FONT_SCALE, FONT_THICKNESS)
|
||
text_w, text_h = text_size
|
||
|
||
# 3. 确定文字位置（避免超出图片范围）
|
||
# 文字默认放在边界框左上角，若左上角空间不足则放在右上角
|
||
text_x = x1 + TEXT_PADDING
|
||
text_y = y1 - TEXT_PADDING - text_h # 文字基线在y轴上方
|
||
if text_y < 0: # 左上角超出图片顶部，调整到右上角
|
||
text_x = x2 - TEXT_PADDING - text_w
|
||
text_y = y1 + TEXT_PADDING + text_h
|
||
|
||
# 4. 绘制文字背景（半透明矩形，避免遮挡图片内容）
|
||
bg_x1 = text_x - TEXT_PADDING
|
||
bg_y1 = text_y - text_h - TEXT_PADDING
|
||
bg_x2 = text_x + text_w + TEXT_PADDING
|
||
bg_y2 = text_y + TEXT_PADDING
|
||
# 确保背景不超出图片范围
|
||
bg_x1 = max(0, bg_x1)
|
||
bg_y1 = max(0, bg_y1)
|
||
bg_x2 = min(img_w, bg_x2)
|
||
bg_y2 = min(img_h, bg_y2)
|
||
|
||
# 半透明背景：先创建背景层，再与原图混合
|
||
bg = img[bg_y1:bg_y2, bg_x1:bg_x2].copy()
|
||
bg = cv2.rectangle(bg, (0, 0), (bg_x2 - bg_x1, bg_y2 - bg_y1), color, -1) # 实心矩形
|
||
img[bg_y1:bg_y2, bg_x1:bg_x2] = cv2.addWeighted(bg, TEXT_BG_OPACITY, img[bg_y1:bg_y2, bg_x1:bg_x2], 1 - TEXT_BG_OPACITY, 0)
|
||
|
||
# 5. 绘制类别名称
|
||
cv2.putText(img, class_name, (text_x, text_y), FONT_FACE, FONT_SCALE, (0, 0, 0), FONT_THICKNESS) # 白色文字
|
||
|
||
return img
|
||
|
||
|
||
# -------------------------- 3. 主函数 --------------------------
|
||
def main():
|
||
# 1. 创建输出目录（若不存在）
|
||
os.makedirs(OUTPUT_IMAGE_DIR, exist_ok=True)
|
||
print(f"标注后的图片将保存到：{OUTPUT_IMAGE_DIR}\n")
|
||
|
||
# 2. 构建类别ID到（名称+颜色）的映射字典
|
||
class_map = {cls_id: (cls_name, cls_color) for cls_id, cls_name, cls_color in CLASS_CONFIG}
|
||
print("类别配置：")
|
||
for cls_id, cls_name, cls_color in CLASS_CONFIG:
|
||
print(f" ID {cls_id} → 名称：{cls_name}，颜色：{cls_color}")
|
||
print()
|
||
|
||
# 3. 获取所有图片文件（仅处理支持的格式）
|
||
image_files = [f for f in os.listdir(INPUT_IMAGE_DIR) if f.lower().endswith(SUPPORTED_IMAGE_FORMATS)]
|
||
if not image_files:
|
||
raise FileNotFoundError(f"在 {INPUT_IMAGE_DIR} 中未找到任何支持的图片文件（{SUPPORTED_IMAGE_FORMATS}）")
|
||
print(f"找到 {len(image_files)} 张图片，开始标注...\n")
|
||
|
||
# 4. 遍历图片并标注
|
||
for img_filename in tqdm(image_files, desc="处理进度"):
|
||
# 4.1 构建图片和标签的路径
|
||
img_name, img_ext = os.path.splitext(img_filename)
|
||
img_path = os.path.join(INPUT_IMAGE_DIR, img_filename)
|
||
label_path = os.path.join(INPUT_LABEL_DIR, f"{img_name}.txt") # 标签文件与图片同名，后缀为txt
|
||
|
||
# 4.2 读取图片（OpenCV默认读取为BGR通道）
|
||
img = cv2.imread(img_path)
|
||
if img is None:
|
||
tqdm.write(f"⚠️ 跳过：无法读取图片 {img_filename}（可能损坏或格式不支持）")
|
||
continue
|
||
img_h, img_w = img.shape[:2]
|
||
|
||
# 4.3 读取标签文件（若不存在则跳过标注，直接保存原图）
|
||
if not os.path.exists(label_path):
|
||
tqdm.write(f"⚠️ 警告：图片 {img_filename} 无对应标签文件 {os.path.basename(label_path)}，直接保存原图")
|
||
annotated_img = img.copy()
|
||
else:
|
||
# 复制原图用于标注（避免修改原始图片）
|
||
annotated_img = img.copy()
|
||
# 读取标签内容
|
||
with open(label_path, 'r', encoding='utf-8') as f:
|
||
label_lines = [line.strip() for line in f.readlines() if line.strip()] # 过滤空行
|
||
|
||
# 4.4 解析每个标签并绘制
|
||
for line_idx, line in enumerate(label_lines):
|
||
try:
|
||
# YOLO标签格式：class_id xc yc w h（空格分隔）
|
||
parts = line.split()
|
||
if len(parts) != 5:
|
||
raise ValueError(f"格式错误（需5个字段，实际{len(parts)}个）")
|
||
|
||
# 解析类别ID和坐标
|
||
cls_id = int(float(parts[0]))
|
||
yolo_coords = [float(p) for p in parts[1:]]
|
||
# 检查YOLO坐标有效性（必须在0-1范围内）
|
||
if not all(0 <= coord <= 1 for coord in yolo_coords):
|
||
raise ValueError(f"YOLO坐标超出0-1范围：{yolo_coords}")
|
||
|
||
# 4.5 转换坐标并绘制
|
||
# 检查类别ID是否在配置中
|
||
if cls_id not in class_map:
|
||
tqdm.write(f"⚠️ 跳过：图片 {img_filename} 标签第{line_idx+1}行，未知类别ID {cls_id}（未在CLASS_CONFIG中配置）")
|
||
continue
|
||
|
||
# 获取类别名称和颜色
|
||
cls_name, cls_color = class_map[cls_id]
|
||
# 转换YOLO坐标为像素坐标
|
||
bbox = yolo2pixel(yolo_coords, img_w, img_h)
|
||
# 绘制标注
|
||
annotated_img = draw_annotation(annotated_img, bbox, cls_name, cls_color)
|
||
|
||
except Exception as e:
|
||
tqdm.write(f"⚠️ 跳过：图片 {img_filename} 标签第{line_idx+1}行解析失败 → {str(e)}")
|
||
continue
|
||
|
||
# 4.6 保存标注后的图片
|
||
output_img_path = os.path.join(OUTPUT_IMAGE_DIR, f"{img_name}_annotated{img_ext}")
|
||
# 保存为JPG格式（若原始是PNG，也可改为img_ext保持原格式）
|
||
# 注：JPG不支持透明通道，若原始是PNG且有透明，建议保留img_ext
|
||
cv2.imwrite(output_img_path, annotated_img)
|
||
|
||
# 5. 完成提示
|
||
print(f"\n✅ 标注完成！共处理 {len(image_files)} 张图片，标注后的图片已保存到：")
|
||
print(f" {OUTPUT_IMAGE_DIR}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
main()
|
||
except Exception as e:
|
||
print(f"\n❌ 程序异常终止：{str(e)}") |