数据增样模版

import json
import os
import random
import cv2
import numpy as np
from tqdm import tqdm

INPUT_DIR = r"E:\20260914_yolov8\images\images"
OUTPUT_DIR = r"E:\20260914_yolov8\images\imgs0_augmented"
os.makedirs(OUTPUT_DIR, exist_ok=True)

---------- 工具函数 ----------

def bbox(points):
xs = [p[0] for p in points]
ys = [p[1] for p in points]
return min(xs), min(ys), max(xs), max(ys)

def points(xmin, ymin, xmax, ymax):
return [[xmin, ymin], [xmax, ymax]]

---------- 1. 翻转 ----------

def flip(img, shapes):
h, w = img.shape[:2]
img = cv2.flip(img, 1) # 水平翻转
for s in shapes:
xmin, ymin, xmax, ymax = bbox(s["points"])
xmin, xmax = w - xmax, w - xmin
s["points"] = points(xmin, ymin, xmax, ymax)
return img, shapes

---------- 2. 旋转 ----------

def rotate(img, shapes, angle):
h, w = img.shape[:2]
M = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1)
img = cv2.warpAffine(img, M, (w, h))

for s in shapes:
xmin, ymin, xmax, ymax = bbox(s["points"])
pts = np.array([[xmin, ymin], [xmax, ymin],
[xmax, ymax], [xmin, ymax]], dtype=np.float32)
pts = cv2.transform(pts.reshape(1, -1, 2), M)[0]

关键：转成 Python float，再取 min/max

xmin = float(max(0, pts[:, 0].min()))
ymin = float(max(0, pts[:, 1].min()))
xmax = float(min(w, pts[:, 0].max()))
ymax = float(min(h, pts[:, 1].max()))

s["points"] = points(xmin, ymin, xmax, ymax)
return img, shapes

---------- 3. 亮度 ----------

def brightness(img, shapes):
img = cv2.convertScaleAbs(img, alpha=1.0, beta=random.randint(-40, 40))
return img, shapes

---------- 4. 加噪 ----------

def noise(img, shapes):
n = np.random.normal(0, 10, img.shape)
img = np.clip(img + n, 0, 255).astype(np.uint8)
return img, shapes

---------- 主流程 ----------

def augment(img, shapes):
"""对一张图做 4 种增强，返回 (后缀, 图像, shapes) 列表"""
results = []

im, sh = flip(img.copy(), json.loads(json.dumps(shapes)))
results.append(("_flip", im, sh))

im, sh = rotate(img.copy(), json.loads(json.dumps(shapes)), 15)
results.append(("_rot", im, sh))

im, sh = brightness(img.copy(), json.loads(json.dumps(shapes)))
results.append(("_bright", im, sh))

im, sh = noise(img.copy(), json.loads(json.dumps(shapes)))
results.append(("_noise", im, sh))

return results

if name == "main":

for name in tqdm(os.listdir(INPUT_DIR)):
if not name.lower().endswith((".jpg", ".png", ".jpeg", ".bmp")):
continue

base = os.path.splitext(name)[0]
img_path = os.path.join(INPUT_DIR, name)
json_path = os.path.join(INPUT_DIR, base + ".json")

if not os.path.exists(json_path):
continue

img = cv2.imread(img_path)
data = json.load(open(json_path, encoding="utf-8"))
shapes = data["shapes"]

for suffix, im, sh in augment(img, shapes):
h, w = im.shape[:2]
new_name = base + suffix

cv2.imwrite(os.path.join(OUTPUT_DIR, new_name + ".jpg"), im)

data["shapes"] = sh
data["imagePath"] = new_name + ".jpg"
data["imageWidth"] = w
data["imageHeight"] = h
json.dump(data, open(os.path.join(OUTPUT_DIR, new_name + ".json"), "w",
encoding="utf-8"), ensure_ascii=False, indent=2)

print("增强完成：", OUTPUT_DIR)