原因是 市面上大多数都是用yolo做 目标检测的,有一些我自己找的实例分割数据增强的代码也都是错的,一类是对图片进行操作了但是没对标签进行操作,另一类是旋转和翻折的标签是错的【主要是这个原因】,其他的没问题。

思路:

        我们手里有的是原始图片(jpg或png)标签(json),且图片和标签的名称是一一对应的,增强过程是需要先将标签的json格式转换成TXT格式,对TXT文件进行代码操作,再将TXT文件转换回json文件格式。  

        往往这个过程都是进行手动操作点来点去,我直接将所有部分都集合成了一个代码。只需要提供图片和标签的位置,再给出生成的位置就可以了。生成的文件夹分别有:图片文件夹、标签文件夹、图片标签文件夹(方便直接在labelme软件中打开查看)。


示例:

只需要修改路径就好(图片和标签在一个路径下也没问题,可以识别):


import os
import cv2
import numpy as np
from pathlib import Path
import shutil
from typing import Tuple, List, Optional
import json
from PIL import Image
from tqdm import tqdm

#  yolo-seg图像增强一步到位代码(可选版)

#  可以在控制台选择要进行的增强方法


class SegmentationAugmenter:
    def __init__(self, images_dir: str, labels_dir: str, output_dir: str, selected_augmentations: List[str]):
        """
        初始化分割数据增强器

        Args:
            images_dir (str): 原始图片目录
            labels_dir (str): 分割标签目录
            output_dir (str): 输出目录
            selected_augmentations (List[str]): 选择的增强方法列表
        """
        self.images_dir = Path(images_dir)
        self.labels_dir = Path(labels_dir)
        self.images_output = Path(output_dir + "_images")
        self.labels_output = Path(output_dir + "_labels")
        self.labelme_output = Path(output_dir + "_labelme")  # 用于存放增强后的图片和json文件
        self.selected_augmentations = selected_augmentations

        # 创建输出目录
        self.images_output.mkdir(parents=True, exist_ok=True)
        self.labels_output.mkdir(parents=True, exist_ok=True)
        self.labelme_output.mkdir(parents=True, exist_ok=True)

        # 支持的图片格式
        self.image_extensions = ('.jpg', '.jpeg', '.png', '.bmp')
        # 标签文件可能的后缀
        self.label_extensions = ('.txt', '.json', '.xml', '.png', '.jpg')

    def adjust_brightness_contrast(self, image: np.ndarray,
                                   brightness_factor: float,
                                   contrast_factor: float) -> np.ndarray:
        """调整图像的亮度和对比度"""
        image = image.astype(np.float32)
        brightness = np.ones(image.shape, dtype=np.float32) * brightness_factor
        bright_img = cv2.add(image, brightness)
        contrast_img = cv2.addWeighted(bright_img, contrast_factor,
                                   np.zeros_like(bright_img), 0, 0)
        return np.clip(contrast_img, 0, 255).astype(np.uint8)

    def add_noise(self, image: np.ndarray, noise_factor: float) -> np.ndarray:
        """添加高斯噪声"""
        image = image.astype(np.float32)
        row, col, ch = image.shape
        noise = np.random.normal(0, noise_factor, (row, col, ch))
        noisy_img = image + noise
        return np.clip(noisy_img, 0, 255).astype(np.uint8)

    def apply_blur(self, image: np.ndarray, kernel_size: int = 5) -> np.ndarray:
        """应用高斯模糊"""
        return cv2.GaussianBlur(image, (kernel_size, kernel_size), 0)

    def apply_sharpen(self, image: np.ndarray) -> np.ndarray:
        """应用锐化"""
        kernel = np.array([[-1, -1, -1],
                           [-1, 9, -1],
                           [-1, -1, -1]])
        return cv2.filter2D(image, -1, kernel)

    def adjust_gamma(self, image: np.ndarray, gamma: float) -> np.ndarray:
        """应用伽马变换"""
        inv_gamma = 1.0 / gamma
        table = np.array([((i / 255.0) ** inv_gamma) * 255
                          for i in np.arange(0, 256)]).astype(np.uint8)
        return cv2.LUT(image, table)

    def adjust_channels(self, image: np.ndarray,
                        b_factor: float, g_factor: float, r_factor: float) -> np.ndarray:
        """调整BGR通道"""
        image = image.astype(np.float32)
        b, g, r = cv2.split(image)
        b = np.clip(b * b_factor, 0, 255)
        g = np.clip(g * g_factor, 0, 255)
        r = np.clip(r * r_factor, 0, 255)
        return cv2.merge([b, g, r]).astype(np.uint8)

    def adjust_hue_saturation(self, image: np.ndarray,
                              hue_factor: float,
                              saturation_factor: float) -> np.ndarray:
        """调整色调和饱和度"""
        hsv_img = cv2.cvtColor(image, cv2.COLOR_BGR2HSV).astype(np.float32)
        hsv_img[:, :, 0] = (hsv_img[:, :, 0] + hue_factor) % 180
        hsv_img[:, :, 1] = np.clip(hsv_img[:, :, 1] * saturation_factor, 0, 255)
        hsv_img = np.clip(hsv_img, 0, 255).astype(np.uint8)
        return cv2.cvtColor(hsv_img, cv2.COLOR_HSV2BGR)

    def find_label_file(self, image_stem: str) -> Optional[Path]:
        """查找对应的标签文件"""
        for ext in self.label_extensions:
            potential_path = self.labels_dir / f"{image_stem}{ext}"
            if potential_path.exists():
                return potential_path
        return None

    def json_to_txt(self, json_path: Path) -> Optional[str]:
        """将json格式的分割标签转换为txt格式"""
        try:
            with open(json_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            txt_content = []
            for shape in data.get('shapes', []):
                label = shape.get('label', '0')
                points = shape.get('points', [])
                if points:
                    # 假设是多边形分割,转换为YOLO格式
                    # YOLO格式: class x1 y1 x2 y2 ... xn yn
                    coords = []
                    for point in points:
                        coords.extend(point)
                    line = f"{label} {' '.join(map(str, coords))}"
                    txt_content.append(line)
            
            return '\n'.join(txt_content)
        except Exception as e:
            print(f"转换json到txt失败 {json_path}: {e}")
            return None
    
    def txt_to_json(self, txt_content: str, image_path: Path, image_data: Optional[np.ndarray] = None) -> Optional[dict]:
        """将txt格式的分割标签转换为json格式"""
        try:
            # 获取图片信息
            if image_data is not None:
                height, width = image_data.shape[:2]
            else:
                # 尝试使用PIL库读取图片
                try:
                    pil_image = Image.open(image_path)
                    image = np.array(pil_image)
                    height, width = image.shape[:2]
                except Exception as e:
                    print(f"读取图片信息失败: {e}")
                    return None
            
            # 解析txt内容
            shapes = []
            lines = txt_content.strip().split('\n')
            for line in lines:
                if not line:
                    continue
                parts = line.split()
                if len(parts) < 3:
                    continue
                
                label = parts[0]
                points = []
                for i in range(1, len(parts), 2):
                    if i + 1 < len(parts):
                        x = float(parts[i])
                        y = float(parts[i+1])
                        points.append([x, y])
                
                if points:
                    shapes.append({
                        "label": label,
                        "points": points,
                        "group_id": None,
                        "shape_type": "polygon",
                        "flags": {}
                    })
            
            # 构建json数据
            json_data = {
                "version": "4.5.6",
                "flags": {},
                "shapes": shapes,
                "imagePath": image_path.name,
                "imageData": None,
                "imageHeight": height,
                "imageWidth": width
            }
            
            return json_data
        except Exception as e:
            print(f"转换txt到json失败: {e}")
            return None
    
    def load_label(self, label_path: Path) -> Optional[str]:
        """加载分割标签并转换为txt格式"""
        if label_path.exists():
            try:
                if label_path.suffix.lower() == '.json':
                    return self.json_to_txt(label_path)
                else:
                    with open(label_path, 'r', encoding='utf-8') as f:
                        return f.read()
            except Exception as e:
                print(f"读取标签文件失败 {label_path}: {e}")
                return None
        return None

    def flip_image(self, image: np.ndarray, flip_code: int) -> np.ndarray:
        """翻转图像
        flip_code: 0=垂直翻转, 1=水平翻转, -1=水平垂直翻转
        """
        return cv2.flip(image, flip_code)
    
    def rotate_image(self, image: np.ndarray, angle: float) -> Tuple[np.ndarray, Tuple[int, int, int, int]]:
        """旋转图像,确保旋转后的图像完全在原始图像范围内,并裁剪掉多余的黑色区域
        angle: 旋转角度
        返回: (裁剪后的图像, 裁剪区域坐标 (x, y, width, height))
        """
        height, width = image.shape[:2]
        
        # 计算旋转后的图像大小,确保所有内容都在边界内
        angle_rad = np.radians(angle)
        cos_theta = abs(np.cos(angle_rad))
        sin_theta = abs(np.sin(angle_rad))
        
        # 计算旋转后的图像大小
        new_width = int((height * sin_theta) + (width * cos_theta))
        new_height = int((height * cos_theta) + (width * sin_theta))
        
        # 计算缩放比例,使旋转后的图像完全包含在原始图像中
        scale = min(width / new_width, height / new_height)
        
        # 计算新的旋转中心
        center = (width / 2, height / 2)
        
        # 创建旋转矩阵
        M = cv2.getRotationMatrix2D(center, angle, scale)
        
        # 计算平移量,使旋转后的图像居中
        tx = (width - new_width * scale) / 2
        ty = (height - new_height * scale) / 2
        M[0, 2] += tx
        M[1, 2] += ty
        
        # 执行旋转
        rotated = cv2.warpAffine(image, M, (width, height), borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0))
        
        # 找到非黑色区域的边界
        gray = cv2.cvtColor(rotated, cv2.COLOR_BGR2GRAY)
        _, thresh = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY)
        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        if contours:
            # 找到最大的轮廓(即图像内容)
            largest_contour = max(contours, key=cv2.contourArea)
            x, y, w, h = cv2.boundingRect(largest_contour)
            
            # 裁剪图像
            cropped = rotated[y:y+h, x:x+w]
            return cropped, (x, y, w, h)
        else:
            # 如果没有找到非黑色区域,返回原始旋转图像
            return rotated, (0, 0, width, height)
    
    def flip_label(self, label_content: str, image_width: int, image_height: int, flip_code: int) -> str:
        """翻转标签
        flip_code: 0=垂直翻转, 1=水平翻转, -1=水平垂直翻转
        """
        if not label_content:
            return label_content
        
        lines = label_content.strip().split('\n')
        flipped_lines = []
        
        for line in lines:
            if not line:
                continue
            parts = line.split()
            if len(parts) < 3:
                flipped_lines.append(line)
                continue
            
            label = parts[0]
            points = []
            for i in range(1, len(parts), 2):
                if i + 1 < len(parts):
                    x = float(parts[i])
                    y = float(parts[i+1])
                    points.append([x, y])
            
            # 翻转点
            flipped_points = []
            for x, y in points:
                if flip_code == 1:  # 水平翻转
                    flipped_x = image_width - x
                    flipped_points.append([flipped_x, y])
                elif flip_code == 0:  # 垂直翻转
                    flipped_y = image_height - y
                    flipped_points.append([x, flipped_y])
                elif flip_code == -1:  # 水平垂直翻转
                    flipped_x = image_width - x
                    flipped_y = image_height - y
                    flipped_points.append([flipped_x, flipped_y])
            
            # 重新构建行
            flipped_parts = [label]
            for x, y in flipped_points:
                flipped_parts.extend([str(x), str(y)])
            flipped_lines.append(' '.join(flipped_parts))
        
        return '\n'.join(flipped_lines)
    
    def line_intersection(self, p1, p2, p3, p4):
        """计算两条线段的交点
        p1, p2: 第一条线段的两个端点
        p3, p4: 第二条线段的两个端点
        返回交点坐标,如果没有交点则返回None
        """
        x1, y1 = p1
        x2, y2 = p2
        x3, y3 = p3
        x4, y4 = p4
        
        denom = (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4)
        if abs(denom) < 1e-10:
            return None
        
        t = ((x1 - x3) * (y3 - y4) - (y1 - y3) * (x3 - x4)) / denom
        
        if t < 0 or t > 1:
            return None
        
        x = x1 + t * (x2 - x1)
        y = y1 + t * (y2 - y1)
        
        return [x, y]
    
    def clip_point_to_boundary(self, x, y, image_width, image_height):
        """将点裁剪到图像边界内,并返回边界上的点
        """
        # 图像边界
        left, right = 0, image_width - 1
        top, bottom = 0, image_height - 1
        
        # 如果点已经在边界内,直接返回
        if left <= x <= right and top <= y <= bottom:
            return [x, y]
        
        # 计算中心点
        center_x = image_width / 2
        center_y = image_height / 2
        
        # 从中心到该点的射线与边界的交点
        dx = x - center_x
        dy = y - center_y
        
        # 避免除以零
        if abs(dx) < 1e-10 and abs(dy) < 1e-10:
            return [center_x, center_y]
        
        # 计算与四条边的交点
        intersections = []
        
        # 左边 x = left
        if abs(dx) > 1e-10:
            t = (left - center_x) / dx
            if t >= 0:
                iy = center_y + t * dy
                if top <= iy <= bottom:
                    intersections.append([left, iy])
        
        # 右边 x = right
        if abs(dx) > 1e-10:
            t = (right - center_x) / dx
            if t >= 0:
                iy = center_y + t * dy
                if top <= iy <= bottom:
                    intersections.append([right, iy])
        
        # 上边 y = top
        if abs(dy) > 1e-10:
            t = (top - center_y) / dy
            if t >= 0:
                ix = center_x + t * dx
                if left <= ix <= right:
                    intersections.append([ix, top])
        
        # 下边 y = bottom
        if abs(dy) > 1e-10:
            t = (bottom - center_y) / dy
            if t >= 0:
                ix = center_x + t * dx
                if left <= ix <= right:
                    intersections.append([ix, bottom])
        
        # 返回最近的交点
        if intersections:
            min_dist = float('inf')
            closest = intersections[0]
            for pt in intersections:
                dist = (pt[0] - x) ** 2 + (pt[1] - y) ** 2
                if dist < min_dist:
                    min_dist = dist
                    closest = pt
            return closest
        
        # 如果没有找到交点,返回边界内的最近点
        return [max(left, min(x, right)), max(top, min(y, bottom))]
    
    def rotate_label(self, label_content: str, image_width: int, image_height: int, angle: float, crop_coords: Tuple[int, int, int, int]) -> str:
        """旋转标签,使用与图像旋转相同的参数,并调整裁剪后的坐标
        angle: 旋转角度
        crop_coords: 裁剪区域坐标 (x, y, width, height)
        """
        if not label_content:
            return label_content
        
        lines = label_content.strip().split('\n')
        rotated_lines = []
        
        # 计算与图像旋转相同的参数
        height, width = image_height, image_width
        angle_rad = np.radians(angle)
        cos_theta = abs(np.cos(angle_rad))
        sin_theta = abs(np.sin(angle_rad))
        
        # 计算旋转后的图像大小
        new_width = int((height * sin_theta) + (width * cos_theta))
        new_height = int((height * cos_theta) + (width * sin_theta))
        
        # 计算缩放比例,使旋转后的图像完全包含在原始图像中
        scale = min(width / new_width, height / new_height)
        
        # 计算新的旋转中心
        center = (width / 2, height / 2)
        
        # 创建与图像旋转相同的旋转矩阵
        M = cv2.getRotationMatrix2D(center, angle, scale)
        
        # 计算平移量,使旋转后的图像居中
        tx = (width - new_width * scale) / 2
        ty = (height - new_height * scale) / 2
        M[0, 2] += tx
        M[1, 2] += ty
        
        # 裁剪坐标
        crop_x, crop_y, crop_w, crop_h = crop_coords
        
        for line in lines:
            if not line:
                continue
            parts = line.split()
            if len(parts) < 3:
                rotated_lines.append(line)
                continue
            
            label = parts[0]
            points = []
            for i in range(1, len(parts), 2):
                if i + 1 < len(parts):
                    x = float(parts[i])
                    y = float(parts[i+1])
                    points.append([x, y])
            
            # 使用与图像旋转相同的旋转矩阵来旋转点
            rotated_points = []
            for i, (x, y) in enumerate(points):
                # 应用旋转矩阵
                rotated_x = M[0, 0] * x + M[0, 1] * y + M[0, 2]
                rotated_y = M[1, 0] * x + M[1, 1] * y + M[1, 2]
                
                # 调整裁剪后的坐标
                adjusted_x = rotated_x - crop_x
                adjusted_y = rotated_y - crop_y
                
                # 限制坐标在裁剪后的图像范围内
                adjusted_x = max(0, min(adjusted_x, crop_w - 1))
                adjusted_y = max(0, min(adjusted_y, crop_h - 1))
                rotated_points.append([adjusted_x, adjusted_y])
            
            # 重新构建行
            rotated_parts = [label]
            for x, y in rotated_points:
                rotated_parts.extend([str(x), str(y)])
            rotated_lines.append(' '.join(rotated_parts))
        
        return '\n'.join(rotated_lines)
    
    def apply_augmentations(self, image: np.ndarray, label_content: Optional[str] = None) -> List[
        Tuple[np.ndarray, Optional[str], str]]:
        """
        应用多种数据增强方法

        对于分割任务,只对图像进行增强,标签内容保持不变
        """
        augmented_data = []
        height, width = image.shape[:2]

        # 1. 亮度调整(+30%)
        if 'bright' in self.selected_augmentations:
            bright_img = self.adjust_brightness_contrast(image, 30, 1.0)  # 亮度+30%
            augmented_data.append((bright_img, label_content, '_bright'))

        # 2. 亮度调整(-30%)
        if 'dark' in self.selected_augmentations:
            dark_img = self.adjust_brightness_contrast(image, -30, 1.0)  # 亮度-30%
            augmented_data.append((dark_img, label_content, '_dark'))

        # 3. 添加噪声(幅度30%)
        if 'noise' in self.selected_augmentations:
            noisy_img = self.add_noise(image, 30)
            augmented_data.append((noisy_img, label_content, '_noise'))

        # 4. 模糊
        if 'blur' in self.selected_augmentations:
            blur_img = self.apply_blur(image, kernel_size=25)  # 更强的模糊
            augmented_data.append((blur_img, label_content, '_blur'))

        # 5. 锐化
        if 'sharp' in self.selected_augmentations:
            sharp_img = self.apply_sharpen(image)
            augmented_data.append((sharp_img, label_content, '_sharp'))

        # 6. 水平翻转
        if 'h_flip' in self.selected_augmentations:
            h_flip = self.flip_image(image, 1)  # 水平翻转
            h_flip_label = self.flip_label(label_content, width, height, 1)
            augmented_data.append((h_flip, h_flip_label, '_h_flip'))

        # 7. 垂直翻转
        if 'v_flip' in self.selected_augmentations:
            v_flip = self.flip_image(image, 0)  # 垂直翻转
            v_flip_label = self.flip_label(label_content, width, height, 0)
            augmented_data.append((v_flip, v_flip_label, '_v_flip'))

        # 8. 旋转45°
        if 'rotate_45' in self.selected_augmentations:
            rotate_45, crop_coords_45 = self.rotate_image(image, 45)
            rotate_45_label = self.rotate_label(label_content, width, height, 45, crop_coords_45)
            augmented_data.append((rotate_45, rotate_45_label, '_rotate_45'))

        # 9. 旋转-45°
        if 'rotate_neg_45' in self.selected_augmentations:
            rotate_neg_45, crop_coords_neg_45 = self.rotate_image(image, -45)
            rotate_neg_45_label = self.rotate_label(label_content, width, height, -45, crop_coords_neg_45)
            augmented_data.append((rotate_neg_45, rotate_neg_45_label, '_rotate_neg_45'))

        return augmented_data

    def process_single_image(self, image_path: Path) -> bool:
        """处理单张图片及其分割标签,返回是否成功"""
        # 读取图片
        try:
            # 使用PIL库读取图片,解决中文字符路径问题
            pil_image = Image.open(image_path)
            # 转换为OpenCV格式
            image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
        except Exception as e:
            print(f"读取图片时出错: {e}")
            return False

        # 查找对应的标签文件
        label_path = self.find_label_file(image_path.stem)

        if label_path is None:
            print(f"找不到标签文件: {image_path.stem}")
            return False

        print(f"找到标签: {label_path.name}")

        # 加载标签内容
        label_content = self.load_label(label_path)
        if label_content is None:
            return False

        # 调整原始图像大小为640x640
        resized_image = cv2.resize(image, (640, 640))
        
        # 调整原始标签坐标
        resized_label_content = label_content
        if label_content is not None:
            # 获取原始图像大小
            orig_height, orig_width = image.shape[:2]
            # 计算缩放比例
            scale_x = 640 / orig_width
            scale_y = 640 / orig_height
            
            # 调整标签坐标
            lines = label_content.strip().split('\n')
            resized_lines = []
            for line in lines:
                if not line:
                    continue
                parts = line.split()
                if len(parts) < 3:
                    resized_lines.append(line)
                    continue
                
                label = parts[0]
                points = []
                for i in range(1, len(parts), 2):
                    if i + 1 < len(parts):
                        x = float(parts[i]) * scale_x
                        y = float(parts[i+1]) * scale_y
                        points.append([x, y])
                
                # 重新构建行
                resized_parts = [label]
                for x, y in points:
                    resized_parts.extend([str(x), str(y)])
                resized_lines.append(' '.join(resized_parts))
            
            resized_label_content = '\n'.join(resized_lines)
        
        # 保存调整后的原始图片
        original_image_path = self.images_output / image_path.name
        try:
            pil_image = Image.fromarray(cv2.cvtColor(resized_image, cv2.COLOR_BGR2RGB))
            pil_image.save(original_image_path)
        except Exception as e:
            print(f"保存原始图片失败: {e}")
        
        # 保存调整后的原始标签
        if resized_label_content is not None:
            original_label_path = self.labels_output / label_path.name
            with open(original_label_path, 'w', encoding='utf-8') as f:
                f.write(resized_label_content)
            
            # 转换为json格式并保存到labelme_output目录
            json_data = self.txt_to_json(resized_label_content, original_image_path, resized_image)
            if json_data:
                json_name = f"{image_path.stem}.json"
                with open(self.labelme_output / json_name, 'w', encoding='utf-8') as f:
                    json.dump(json_data, f, ensure_ascii=False, indent=2)
                # 保存调整后的图片到labelme_output目录
                shutil.copy2(original_image_path, self.labelme_output / image_path.name)

        # 应用增强
        augmented_data = self.apply_augmentations(image, label_content)

        # 保存增强后的图片和对应的标签
        for aug_image, aug_label_content, suffix in augmented_data:
            # 调整图像大小为640x640
            resized_image = cv2.resize(aug_image, (640, 640))
            
            # 调整标签坐标
            resized_label_content = aug_label_content
            if aug_label_content is not None:
                # 获取原始图像大小
                orig_height, orig_width = aug_image.shape[:2]
                # 计算缩放比例
                scale_x = 640 / orig_width
                scale_y = 640 / orig_height
                
                # 调整标签坐标
                lines = aug_label_content.strip().split('\n')
                resized_lines = []
                for line in lines:
                    if not line:
                        continue
                    parts = line.split()
                    if len(parts) < 3:
                        resized_lines.append(line)
                        continue
                    
                    label = parts[0]
                    points = []
                    for i in range(1, len(parts), 2):
                        if i + 1 < len(parts):
                            x = float(parts[i]) * scale_x
                            y = float(parts[i+1]) * scale_y
                            points.append([x, y])
                    
                    # 重新构建行
                    resized_parts = [label]
                    for x, y in points:
                        resized_parts.extend([str(x), str(y)])
                    resized_lines.append(' '.join(resized_parts))
                
                resized_label_content = '\n'.join(resized_lines)
            
            # 保存增强后的图片
            aug_image_name = f"{image_path.stem}{suffix}{image_path.suffix}"
            aug_image_path = self.images_output / aug_image_name
            # 使用PIL库保存图片,解决中文字符路径问题
            try:
                # 转换为PIL格式并保存
                pil_image = Image.fromarray(cv2.cvtColor(resized_image, cv2.COLOR_BGR2RGB))
                pil_image.save(aug_image_path)
            except Exception as e:
                print(f"保存图片失败: {e}")
                continue

            # 保存对应的标签(调整后的内容)
            if resized_label_content is not None:
                aug_label_name = f"{image_path.stem}{suffix}{label_path.suffix}"
                with open(self.labels_output / aug_label_name, 'w', encoding='utf-8') as f:
                    f.write(resized_label_content)
                
                # 转换为json格式并保存到labelme_output目录
                json_data = self.txt_to_json(resized_label_content, aug_image_path, resized_image)
                if json_data:
                    aug_json_name = f"{image_path.stem}{suffix}.json"
                    with open(self.labelme_output / aug_json_name, 'w', encoding='utf-8') as f:
                        json.dump(json_data, f, ensure_ascii=False, indent=2)
                    shutil.copy2(aug_image_path, self.labelme_output / aug_image_name)

        return True

    def process_dataset(self) -> None:
        """处理整个数据集"""
        image_files = [f for f in self.images_dir.iterdir()
                       if f.suffix.lower() in self.image_extensions]

        if not image_files:
            print(f"在 {self.images_dir} 中没有找到图片文件")
            return

        total_images = len(image_files)
        successful_count = 0

        print(f"开始处理数据集...")
        print(f"共找到 {total_images} 个原始图片")
        print(f"图片目录: {self.images_dir}")
        print(f"标签目录: {self.labels_dir}")
        print(f"选择的增强方法: {', '.join(self.selected_augmentations)}")

        # 计算预计产出照片数量
        num_selected_augmentations = len(self.selected_augmentations)
        expected_total = total_images * (1 + num_selected_augmentations)  # 原始图片 + 增强后的图片
        print(f"预计产出照片数量: {expected_total}")

        # 使用tqdm添加进度条
        for idx, image_path in enumerate(tqdm(image_files, desc="处理图片", unit="张"), 1):
            if self.process_single_image(image_path):
                successful_count += 1

        # 统计结果
        final_images = len(list(self.images_output.glob('*')))
        final_labels = len(list(self.labels_output.glob('*')))
        final_labelme_files = len(list(self.labelme_output.glob('*'))) // 2  # 每个图片对应一个json文件
        print(f"\n增强完成!")
        print(f"原始图片数量: {total_images}")
        print(f"成功处理的图片: {successful_count}")
        print(f"增强后图片总数: {final_images}")
        print(f"增强后标签总数: {final_labels}")
        print(f"增强后可用于labelme查看的文件对数: {final_labelme_files}")
        print(f"增强后的图片保存在: {self.images_output}")
        print(f"增强后的标签保存在: {self.labels_output}")
        print(f"增强后的图片和json文件保存在: {self.labelme_output}")


def inspect_label_file(label_path: str):
    """检查标签文件内容"""
    try:
        with open(label_path, 'r', encoding='utf-8') as f:
            content = f.read()
        print(f"\n标签文件内容示例 ({os.path.basename(label_path)}):")
        print("前500个字符:")
        print(content[:500])
        print(f"文件总长度: {len(content)} 字符")
        return True
    except Exception as e:
        print(f"检查标签文件失败: {e}")
        return False


def main():
    # 设置路径
    images_dir = "图片路径"
    labels_dir = "标签路径"
    output_dir = "结果路径"

    # 显示可选择的增强方法
    print("\n===== 数据增强方法选择 =====")
    print("请选择要进行的增强方法,输入对应的数字(多个数字用空格分隔):")
    print("1. 亮度+30% (bright)")
    print("2. 亮度-30% (dark)")
    print("3. 添加噪声 (noise)")
    print("4. 模糊 (blur)")
    print("5. 锐化 (sharp)")
    print("6. 水平翻转 (h_flip)")
    print("7. 垂直翻转 (v_flip)")
    print("8. 旋转45° (rotate_45)")
    print("9. 旋转-45° (rotate_neg_45)")
    
    # 映射数字到增强方法
    aug_map = {
        '1': 'bright',
        '2': 'dark',
        '3': 'noise',
        '4': 'blur',
        '5': 'sharp',
        '6': 'h_flip',
        '7': 'v_flip',
        '8': 'rotate_45',
        '9': 'rotate_neg_45'
    }
    
    # 获取用户输入
    user_input = input("请输入选择的数字: ")
    selected_nums = user_input.strip().split()
    
    # 验证输入并转换为增强方法
    selected_augmentations = []
    for num in selected_nums:
        if num in aug_map:
            selected_augmentations.append(aug_map[num])
        else:
            print(f"警告: 无效的选择 {num},将被忽略")
    
    # 检查是否有选择
    if not selected_augmentations:
        print("错误: 请至少选择一种增强方法")
        return
    
    # 检查目录和标签文件内容
    print("\n检查目录和标签文件...")
    print(f"图片目录存在: {os.path.exists(images_dir)}")
    print(f"标签目录存在: {os.path.exists(labels_dir)}")

    # 检查一个标签文件的内容
    label_files = list(Path(labels_dir).glob("*.txt"))
    if not label_files:
        label_files = list(Path(labels_dir).glob("*.json"))
    
    if label_files:
        print(f"\n检查第一个标签文件内容:")
        inspect_label_file(str(label_files[0]))

    # 创建增强器并处理数据集
    augmenter = SegmentationAugmenter(images_dir, labels_dir, output_dir, selected_augmentations)
    augmenter.process_dataset()


if __name__ == "__main__":
    main()

Logo

腾讯云面向开发者汇聚海量精品云计算使用和开发经验,营造开放的云计算技术生态圈。

更多推荐