Spaces:
Sleeping
Sleeping
File size: 2,373 Bytes
9cf599c | 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 | import os
from PIL import Image, ImageOps, ImageEnhance
import random
# Directory containing normal images for training
NORMAL_DIR = 'dataset/train/normal'
TARGET_COUNT = 100
# List all images in the directory
images = [f for f in os.listdir(NORMAL_DIR) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
current_count = len(images)
print(f"[INFO] Found {current_count} images in {NORMAL_DIR}")
if current_count == 0:
print("[ERROR] No images found to augment.")
exit(1)
# Augmentation functions
AUGS = [
lambda img: img.rotate(random.randint(-30, 30)),
lambda img: ImageOps.mirror(img),
lambda img: ImageOps.flip(img),
lambda img: ImageEnhance.Brightness(img).enhance(random.uniform(0.7, 1.3)),
lambda img: ImageEnhance.Contrast(img).enhance(random.uniform(0.7, 1.3)),
lambda img: ImageEnhance.Color(img).enhance(random.uniform(0.7, 1.3)),
]
aug_idx = 0
while current_count < TARGET_COUNT:
for fname in images:
if current_count >= TARGET_COUNT:
break
img_path = os.path.join(NORMAL_DIR, fname)
img = Image.open(img_path).convert('RGB')
aug = random.choice(AUGS)
aug_img = aug(img)
aug_fname = f"aug_{aug_idx}_{fname}"
aug_img.save(os.path.join(NORMAL_DIR, aug_fname))
aug_idx += 1
current_count += 1
print(f"[AUG] Saved {aug_fname}")
print(f"[DONE] {NORMAL_DIR} now contains {current_count} images.")
# Repeat for faulty images in dataset/test/faulty
FAULTY_DIR = 'dataset/test/faulty'
TARGET_COUNT = 100
images = [f for f in os.listdir(FAULTY_DIR) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
current_count = len(images)
print(f"[INFO] Found {current_count} images in {FAULTY_DIR}")
if current_count == 0:
print("[ERROR] No faulty images found to augment.")
exit(1)
aug_idx = 0
while current_count < TARGET_COUNT:
for fname in images:
if current_count >= TARGET_COUNT:
break
img_path = os.path.join(FAULTY_DIR, fname)
img = Image.open(img_path).convert('RGB')
aug = random.choice(AUGS)
aug_img = aug(img)
aug_fname = f"aug_{aug_idx}_{fname}"
aug_img.save(os.path.join(FAULTY_DIR, aug_fname))
aug_idx += 1
current_count += 1
print(f"[AUG] Saved {aug_fname}")
print(f"[DONE] {FAULTY_DIR} now contains {current_count} images.")
|