Spaces:
Sleeping
Sleeping
File size: 1,375 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 | import os
import shutil
import random
TX_DIR = 'TX'
TRAIN_NORMAL = 'dataset/train/normal'
TRAIN_FAULTY = 'dataset/train/faulty'
TARGET_COUNT = 100
# Clean train folders
shutil.rmtree(TRAIN_NORMAL, ignore_errors=True)
os.makedirs(TRAIN_NORMAL, exist_ok=True)
shutil.rmtree(TRAIN_FAULTY, ignore_errors=True)
os.makedirs(TRAIN_FAULTY, exist_ok=True)
# Helper to collect images from TX folders
def collect_images(src_pattern, dst_folder, needed):
collected = 0
for root, dirs, files in os.walk(TX_DIR):
if src_pattern in root:
imgs = [f for f in files if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
random.shuffle(imgs)
for f in imgs:
if collected >= needed:
return
src_path = os.path.join(root, f)
dst_path = os.path.join(dst_folder, f)
if not os.path.exists(dst_path):
shutil.copyfile(src_path, dst_path)
collected += 1
print(f"[COPY] {src_path} -> {dst_path}")
# Fill train/normal and train/faulty to TARGET_COUNT from TX
collect_images('normal', TRAIN_NORMAL, TARGET_COUNT)
collect_images('faulty', 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))}")
|