图像分割离线数据增强(一步到位版)
·
原因是 市面上大多数都是用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()
更多推荐
所有评论(0)