create_db.py

import os
import difflib
import csv

# Define the base directory
base_dir = r"C:\Users\bradl\Documents\GitHub\occlusion-study\mocap-notebooks-opensim"

# Define the subdirectories to compare
# sub_dirs = ['1', '2', '3']  # Add all your subdirectories here
sub_dirs = ['5']  # Add all your subdirectories here
views = ['front', 'sag']


def find_most_similar_filename(target_file, file_list):
    """Find the most similar filename in file_list to the target_file."""
    matches = difflib.get_close_matches(target_file, file_list, n=1, cutoff=0.0)
    return matches[0] if matches else None


def main():
    total_pairs = 0
    csv_file = 'similar_files.csv'

    with open(csv_file, mode='w', newline='') as file:
        writer = csv.writer(file)
        writer.writerow(['Kinect File', 'MM File', 'StartK', 'StopK', 'CorrAngle', 'Process'])

        for sub_dir in sub_dirs:
            for view in views:
                kinect_path = os.path.join(base_dir, sub_dir, 'kinect', view)
                mm_path = os.path.join(base_dir, sub_dir, 'mm', view)

                if not os.path.exists(kinect_path) or not os.path.exists(mm_path):
                    print(f"Paths do not exist: {kinect_path} or {mm_path}")
                    continue

                kinect_files = [f for f in os.listdir(kinect_path) if f.endswith('.trc') and '_rot' not in f and '_cut' not in f]
                mm_files = [f for f in os.listdir(mm_path) if f.endswith('.trc') and '_rot' not in f and '_cut' not in f]

                for kinect_file in kinect_files:
                    most_similar_mm_file = find_most_similar_filename(kinect_file, mm_files)
                    if most_similar_mm_file:
                        total_pairs += 1
                        kinect_file_path = os.path.join(kinect_path, kinect_file)
                        mm_file_path = os.path.join(mm_path, most_similar_mm_file)
                        start_k = 0
                        stop_k = 0
                        CorrAngle = 0
                        Process = 1

                        writer.writerow([kinect_file_path, mm_file_path, start_k, stop_k, CorrAngle, Process])
                        print(f"Kinect file: {kinect_file_path}")
                        print(f"MM file: {mm_file_path}")
                        print("------")
                    else:
                        print(f"No similar file found for {kinect_file} in {mm_path}")

    print(f"Total pairs found: {total_pairs}")


if __name__ == "__main__":
    main()