| 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() |
