Spaces:
Sleeping
Sleeping
| 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))}") | |