用python实现根据yolov5边框标签文件黑化边框外背景

原图

标签

处理后

代码

# -*- coding: utf-8 -*-
import os
import cv2
import numpy as np

# 配置路径（请根据实际情况修改）
input_img_dir = r'***'
input_label_dir = r'***'
output_dir = r'***'

# 创建输出目录
os.makedirs(output_dir, exist_ok=True)

def process_images():
    # 遍历图片目录
    for img_file in os.listdir(input_img_dir):
        # 构建完整文件路径
        img_path = os.path.join(input_img_dir, img_file)
        
        # 只处理图片文件
        if not img_file.lower().endswith(('.png', '.jpg', '.jpeg')):
            continue

        # 读取对应标签文件
        label_file = os.path.splitext(img_file)[0] + '.txt'
        label_path = os.path.join(input_label_dir, label_file)

        if not os.path.exists(label_path):
            print(f"警告：找不到 {img_file} 对应的标签文件，已跳过")
            continue

        # 读取图片
        # img = cv2.imread(img_path)
        # if img is None:
        #     print(f"错误：无法读取图片 {img_file}，已跳过")
        #     continue
        
        # 读取图片
        try:
            # 使用二进制模式打开图片文件
            with open(img_path, 'rb') as f:
                # 将图片数据读入numpy数组
                chunk = f.read()
                img_array = np.frombuffer(chunk, dtype=np.uint8)
                # 解码图片
                img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
            if img is None:
                raise ValueError(f"无法解码图片 {img_file}")
        except Exception as e:
            print(f"错误：{e} | 文件：{img_file}")
            continue

        # 创建全黑背景
        height, width = img.shape[:2] # 获取图像的高度(height)和宽度(width)两个维度信息
        masked = np.zeros_like(img) # 创建黑色画布，如果要白色画布可以masked = np.full_like(img, 255)

        # 读取并解析YOLO标签文件
        with open(label_path, 'r') as f:  # 使用with语句确保文件安全关闭
            for line in f:  # 逐行迭代（内存友好）
                # 预处理行内容
                cleaned_line = line.strip()
                
                # 跳过空行和注释
                if not cleaned_line or cleaned_line.startswith('#'):
                    continue
                
                # 数据分割与验证
                parts = cleaned_line.split()
                if len(parts) < 5:
                    print(f"警告：无效数据行（字段不足），已跳过 - {cleaned_line}")
                    continue
                
                # 数据类型转换（带异常捕获）
                try:
                    data = [float(x) for x in parts]
                except ValueError as e:
                    print(f"数值转换失败：{cleaned_line} | 错误：{str(e)}")
                    continue
                
                # 坐标参数校验（YOLO格式应为归一化0-1）
                if not all(0 <= x <= 1 for x in data[1:5]):
                    print(f"异常坐标值（超出0-1范围），已跳过 - {data[1:]}")
                    continue
                
                # 解析坐标参数
                _, x_norm, y_norm, w_norm, h_norm = data[:5]
                
                # 转换为像素坐标（保持浮点精度计算）
                x_center = x_norm * width
                y_center = y_norm * height
                box_width = w_norm * width
                box_height = h_norm * height
                
                # 计算边界框坐标（四舍五入更精确）
                x1 = int(round(x_center - box_width / 2))
                y1 = int(round(y_center - box_height / 2))
                x2 = int(round(x_center + box_width / 2))
                x2 = min(x2, width)  # 优先约束右边界
                y2 = int(round(y_center + box_height / 2))
                y2 = min(y2, height)  # 优先约束下边界
                
                # 二次校验区域有效性
                if x2 <= x1 or y2 <= y1:
                    print(f"无效区域：x=[{x1}-{x2}], y=[{y1}-{y2}]，已跳过")
                    continue
                
                # 最终安全裁剪（防御性编程）
                safe_x1 = max(0, x1)
                safe_y1 = max(0, y1)
                safe_x2 = min(width, x2)
                safe_y2 = min(height, y2)
                
                # 复制有效区域
                try:
                    masked[safe_y1:safe_y2, safe_x1:safe_x2] = img[safe_y1:safe_y2, safe_x1:safe_x2]
                except ValueError as e:
                    print(f"区域复制失败：{e} | 坐标：x={safe_x1}-{safe_x2}, y={safe_y1}-{safe_y2}")

        # 保存结果
        # output_path = os.path.join(output_dir, img_file)
        # cv2.imwrite(output_path, masked)
        # print(f"已处理：{img_file}")
        
        # 保存结果
        try:
            # 分解文件名和扩展名
            file_name, ext = os.path.splitext(img_file)
            ext = ext.lower()  # 统一转换为小写

            # 设置编码参数
            if ext in ('.jpg', '.jpeg'):
                encode_params = [int(cv2.IMWRITE_JPEG_QUALITY), 95]
            elif ext == '.png':
                encode_params = [int(cv2.IMWRITE_PNG_COMPRESSION), 4]
            else:
                encode_params = []
            
            # 编码并保存
            success, encoded_image = cv2.imencode(ext, masked, encode_params)
            if success:
                output_path = os.path.join(output_dir, img_file)
                with open(output_path, 'wb') as f:
                    encoded_image.tofile(f)
                print(f"已处理：{img_file}")
            else:
                raise ValueError("图像编码失败")
        except Exception as e:
            print(f"错误：保存文件失败 {e} | 文件：{img_file}")

if __name__ == '__main__':
    process_images()
    print("处理完成！所有结果已保存到：", output_dir)