visualize.py

#!/usr/bin/env python3
"""
Fully customizable interactive visualization for IMU sensor data.
Each plot can independently select sensor type AND pose.
Can add more plots dynamically. Refreshing resets to 4 plots.
"""

import pandas as pd
import json
from pathlib import Path
import webbrowser


# Load configuration from config.json
def load_config():
    """Load configuration from config.json."""
    config_path = Path(__file__).parent / "config.json"
    try:
        with open(config_path, 'r') as f:
            return json.load(f)
    except FileNotFoundError:
        raise FileNotFoundError(
            f"Configuration file not found: {config_path}\n"
            "Please ensure config.json exists in the project root."
        )


# Load config at module level
CONFIG = load_config()


def load_all_data(base_path):
    """Load all pose data and prepare it for JavaScript."""
    base_path = Path(base_path)
    pose_prefix = CONFIG['paths']['pose_folder_prefix']
    csv_pattern = CONFIG['paths']['csv_file_pattern']
    pose_folders = sorted([f for f in base_path.iterdir() if f.is_dir() and f.name.startswith(pose_prefix)])
    
    all_data = {}
    
    for pose_folder in pose_folders:
        csv_files = sorted(list(pose_folder.glob(csv_pattern)))
        if csv_files:
            df = pd.read_csv(csv_files[0])
            pose_name = pose_folder.name
            
            # Extract all sensor data based on config
            all_data[pose_name] = {}
            for sensor_name, sensor_config in CONFIG['sensors'].items():
                columns = sensor_config['columns']
                all_data[pose_name][sensor_name] = {
                    'x': df[columns['x']].tolist(),
                    'y': df[columns['y']].tolist(),
                    'z': df[columns['z']].tolist()
                }
    
    return all_data


def save_data_json(all_data, output_file):
    """Save sensor data as JSON file."""
    with open(output_file, 'w') as f:
        json.dump(all_data, f, indent=2)
    
    print(f"✓ Created: {output_file.name}")


def create_data_script(all_data, output_file):
    """Create a JavaScript file with embedded data."""
    data_json = json.dumps(all_data)
    
    js_content = f"""// Sensor data loaded from CSV files
window.SENSOR_DATA = {data_json};
"""
    
    with open(output_file, 'w') as f:
        f.write(js_content)
    
    print(f"✓ Created: {output_file.name}")


def create_config_script(output_file):
    """Create a JavaScript file with embedded configuration."""
    config_json = json.dumps(CONFIG, indent=2)
    
    js_content = f"""// Configuration loaded from config.json
window.CONFIG = {config_json};
"""
    
    with open(output_file, 'w') as f:
        f.write(js_content)
    
    print(f"✓ Created: {output_file.name}")


def main():
    """Main function."""
    data_dir = CONFIG['paths']['data_dir']
    output_dir_name = CONFIG['paths']['output_dir']
    html_filename = CONFIG['paths']['html_file']
    
    base_path = Path(__file__).parent / data_dir
    
    # Check if data folder exists
    if not base_path.exists():
        raise FileNotFoundError(f"Data folder not found: {base_path}\nPlease create the '{data_dir}' folder with pose subdirectories.")
    
    output_dir = Path(__file__).parent / output_dir_name
    output_dir.mkdir(exist_ok=True)
    
    data_file = output_dir / "data.json"
    data_script = output_dir / "data.js"
    config_script = output_dir / "config.js"
    html_file = Path(__file__).parent / html_filename
    
    print("=" * 80)
    print("IMU Sensor Visualization")
    print("=" * 80)
    print("\nFeatures:")
    print("  ✓ Starts with 4 plots in a 2×2 grid")
    print("  ✓ ➕ Add Plot button to add more plots dynamically")
    print("  ✓ ✕ Remove button on each plot (when > 1 plot)")
    print("  ✓ Dropdown menus to select SENSOR TYPE for each plot")
    print("  ✓ Dropdown menus to select POSE for each plot")
    print("  ✓ Refresh page to reset to 4 plots")
    print("  ✓ Full animation control with play/pause/slider")
    print("=" * 80)
    
    print("\nLoading sensor data...")
    all_data = load_all_data(base_path)
    
    print("Creating data files...")
    save_data_json(all_data, data_file)
    create_data_script(all_data, data_script)
    create_config_script(config_script)
    
    print("\n" + "=" * 80)
    print("✅ Visualization ready!")
    print(f"\nData: {data_file}")
    print(f"Data Script: {data_script}")
    print(f"HTML: {html_file}")
    
    if CONFIG['visualization']['auto_open_browser']:
        print("\nOpening in browser...")
        webbrowser.open(f'file://{html_file.absolute()}')
    else:
        print(f"\nOpen in browser: file://{html_file.absolute()}")
    
    print("=" * 80)


if __name__ == "__main__":
    main()