File size: 4,720 Bytes
b42373a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | import os
import cv2
import json
import numpy as np
import albumentations as A
from barcode_scanner import scan_all_barcodes
BARCODE_DIR = "images/barcode"
CHASSIS_DIR = "images/chassis"
OUTPUT_DIR = "images/chassis_augmented"
GT_PATH = "ground_truth.json"
AUGMENTS_PER_IMAGE = 25
def get_augmentation_pipeline():
return A.Compose([
A.OneOf([
A.RandomBrightnessContrast(
brightness_limit=0.4,
contrast_limit=0.4,
p=1.0
),
A.RandomGamma(gamma_limit=(60, 140), p=1.0),
A.CLAHE(clip_limit=4.0, p=1.0),
], p=0.9),
A.OneOf([
A.RandomShadow(
shadow_roi=(0, 0, 1, 1),
num_shadows_lower=1,
num_shadows_upper=2,
shadow_dimension=4,
p=1.0
),
A.RandomSunFlare(
flare_roi=(0, 0, 1, 0.5),
angle_lower=0,
src_radius=80,
p=1.0
),
], p=0.5),
A.OneOf([
A.MotionBlur(blur_limit=(3, 7), p=1.0),
A.GaussianBlur(blur_limit=(3, 5), p=1.0),
A.MedianBlur(blur_limit=3, p=1.0),
], p=0.4),
A.OneOf([
A.GaussNoise(var_limit=(10, 50), p=1.0),
A.ISONoise(color_shift=(0.01, 0.05), intensity=(0.1, 0.5), p=1.0),
A.MultiplicativeNoise(multiplier=(0.9, 1.1), p=1.0),
], p=0.6),
A.OneOf([
A.Perspective(scale=(0.02, 0.08), p=1.0),
A.ShiftScaleRotate(
shift_limit=0.05,
scale_limit=0.1,
rotate_limit=10,
border_mode=cv2.BORDER_REPLICATE,
p=1.0
),
A.ElasticTransform(
alpha=30,
sigma=5,
alpha_affine=5,
border_mode=cv2.BORDER_REPLICATE,
p=1.0
),
], p=0.7),
A.OneOf([
A.ImageCompression(quality_lower=60, quality_upper=95, p=1.0),
A.Downscale(scale_min=0.5, scale_max=0.9, p=1.0),
], p=0.3),
A.OneOf([
A.CoarseDropout(
max_holes=8,
max_height=2,
max_width=30,
min_holes=2,
fill_value=128,
p=1.0
),
A.GridDistortion(num_steps=5, distort_limit=0.1, p=1.0),
], p=0.4),
])
def augment_dataset():
os.makedirs(OUTPUT_DIR, exist_ok=True)
print("[1/3] Loading ground truth from barcodes...")
ground_truths = scan_all_barcodes(BARCODE_DIR)
ground_truths = {k: v for k, v in ground_truths.items() if v}
print(f" Got {len(ground_truths)} labeled pairs")
chassis_files = sorted([
f for f in os.listdir(CHASSIS_DIR)
if f.lower().endswith(('.jpg', '.jpeg', '.png'))
and os.path.splitext(f)[0] in ground_truths
])
print(f" Found {len(chassis_files)} chassis images with labels")
pipeline = get_augmentation_pipeline()
augmented_gt = {}
total = 0
print(f"\n[2/3] Augmenting — {AUGMENTS_PER_IMAGE} variations per image...")
for fname in chassis_files:
key = os.path.splitext(fname)[0]
label = ground_truths[key]
img_path = os.path.join(CHASSIS_DIR, fname)
img = cv2.imread(img_path)
if img is None:
print(f" [SKIP] Could not read {fname}")
continue
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
orig_name = f"{key}_orig.jpg"
cv2.imwrite(os.path.join(OUTPUT_DIR, orig_name), img)
augmented_gt[orig_name] = label
total += 1
for i in range(AUGMENTS_PER_IMAGE):
try:
augmented = pipeline(image=img_rgb)["image"]
aug_bgr = cv2.cvtColor(augmented, cv2.COLOR_RGB2BGR)
aug_name = f"{key}_aug{i:03d}.jpg"
cv2.imwrite(os.path.join(OUTPUT_DIR, aug_name), aug_bgr)
augmented_gt[aug_name] = label
total += 1
except Exception as e:
print(f" [WARN] Augmentation failed for {fname} variation {i}: {e}")
print(f" {key} -> {AUGMENTS_PER_IMAGE + 1} images (label: {label})")
with open(GT_PATH, "w") as f:
json.dump(augmented_gt, f, indent=2)
print(f"\n[3/3] Done!")
print(f" Total images generated : {total}")
print(f" Saved to : {OUTPUT_DIR}/")
print(f" Ground truth saved to : {GT_PATH}")
print(f"\nNext step: use these images to fine-tune PaddleOCR")
if __name__ == "__main__":
augment_dataset() |