change_names.py

import glob
import os
from tqdm import tqdm
import multiprocessing as mp
from functools import partial
import shutil

def rename_files_worker(args):
    """Worker function to rename files for a single item"""
    i, f, dataset_dir, start_index = args
    
    path_to_image = f"{dataset_dir}/{f}_final.png"
    path_to_data = f"{dataset_dir}/{f}_data.pt"
    path_to_metadata = f"{dataset_dir}/{f}_metadata.json"

    new_path_to_image = f"{dataset_dir}/{i + start_index:06d}_final.png"
    new_path_to_data = f"{dataset_dir}/{i + start_index:06d}_data.pt"
    new_path_to_metadata = f"{dataset_dir}/{i + start_index:06d}_metadata.json"

    # Check if files exist before renaming to avoid errors
    if os.path.exists(path_to_image):
        os.rename(path_to_image, new_path_to_image)
    if os.path.exists(path_to_data):
        os.rename(path_to_data, new_path_to_data)
    if os.path.exists(path_to_metadata):
        os.rename(path_to_metadata, new_path_to_metadata)
    
    return f"Processed {f}"

def change_names(dataset_dir, start_index=0, num_processes=None):
    if num_processes is None:
        num_processes = mp.cpu_count()

    fname_list = []

    for f in sorted(glob.glob(os.path.join(dataset_dir, "*.pt"))):
        fname_list.append(f.split("/")[-1].split("_")[0])

    # Prepare arguments for worker processes
    worker_args = [(i, f, dataset_dir, start_index) for i, f in enumerate(fname_list)]
    
    # Use multiprocessing pool with tqdm for progress tracking
    with mp.Pool(processes=num_processes) as pool:
        list(tqdm(
            pool.imap(rename_files_worker, worker_args),
            total=len(worker_args),
            desc=f"Processing files with {num_processes} processes"
        ))

def move_files(source, target):
    for f in tqdm(glob.glob(f"{source}/*"), desc="Moving files"):
        shutil.move(f, f.replace(source, target))

if __name__ == "__main__":
    move_files("/data/inversion_data/0712_image_data", "/data/inversion_data/diffusionDB_enhanced_150k_image_data/")