Spaces:
Sleeping
Sleeping
File size: 2,512 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 | import os
from PIL import Image, ImageOps, ImageEnhance
import random
import shutil
TRAIN_NORMAL = 'dataset/train/normal'
TRAIN_FAULTY = 'dataset/train/faulty'
TEST_NORMAL = 'dataset/test/normal'
TEST_FAULTY = 'dataset/test/faulty'
TARGET_COUNT = 100
os.makedirs(TRAIN_NORMAL, exist_ok=True)
os.makedirs(TRAIN_FAULTY, exist_ok=True)
# Helper: move images from test to train if needed
def move_images(src, dst, needed):
imgs = [f for f in os.listdir(src) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
moved = 0
for f in imgs:
if moved >= needed:
break
shutil.move(os.path.join(src, f), os.path.join(dst, f))
moved += 1
return moved
# 1. Move images from test to train if train folders have < TARGET_COUNT
normal_needed = TARGET_COUNT - len([f for f in os.listdir(TRAIN_NORMAL) if f.lower().endswith(('.jpg', '.jpeg', '.png'))])
faulty_needed = TARGET_COUNT - len([f for f in os.listdir(TRAIN_FAULTY) if f.lower().endswith(('.jpg', '.jpeg', '.png'))])
if normal_needed > 0:
move_images(TEST_NORMAL, TRAIN_NORMAL, normal_needed)
if faulty_needed > 0:
move_images(TEST_FAULTY, TRAIN_FAULTY, faulty_needed)
# 2. Augment if still not enough
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)),
]
def augment_to_count(folder, target):
images = [f for f in os.listdir(folder) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
current_count = len(images)
aug_idx = 0
while current_count < target:
for fname in images:
if current_count >= target:
break
img_path = os.path.join(folder, 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(folder, aug_fname))
aug_idx += 1
current_count += 1
print(f"[AUG] Saved {aug_fname} in {folder}")
augment_to_count(TRAIN_NORMAL, TARGET_COUNT)
augment_to_count(TRAIN_FAULTY, TARGET_COUNT)
print(f"[DONE] {TRAIN_NORMAL} images: {len(os.listdir(TRAIN_NORMAL))}")
print(f"[DONE] {TRAIN_FAULTY} images: {len(os.listdir(TRAIN_FAULTY))}")
|