app.py

from flask import Flask, render_template, request, jsonify, send_file
import numpy as np
import os
import io

app = Flask(__name__)

# Ensure static folder exists for URDFs
if not os.path.exists('static'):
    os.makedirs('static')

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/upload_motion', methods=['POST'])
def upload_motion():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    
    file = request.files['file']
    
    try:
        # Load the NPZ content directly from memory
        data = np.load(file)
        
        # --- KEY UPDATES HERE ---
        # 1. We now look for 'base_pos_w' and 'base_quat_w'
        # 2. We convert standard numpy types to Python lists for JSON serialization
        
        response_data = {
            'joint_names': data['joint_names'].tolist(),
            'joint_pos': data['joint_pos'].tolist(),
            
            # New Field Names
            'base_pos_w': data['base_pos_w'].tolist(),
            'base_quat_w': data['base_quat_w'].tolist(), 
            
            'num_frames': len(data['joint_pos']),
            'framerate': float(data['framerate']) if 'framerate' in data else 30.0
        }
        
        return jsonify(response_data)

    except Exception as e:
        # This catches the "KeyError" if the file doesn't have the right names
        return jsonify({'error': str(e)}), 500
    
@app.route('/save_motion', methods=['POST'])
def save_motion():
    try:
        # 1. Get JSON data from the frontend
        data = request.json
        
        if not data:
            return jsonify({"error": "No data received"}), 400

        # 2. Prepare dictionaries for np.savez
        # We assume the keys in JSON match the original npz keys
        output_dict = {}

        # Helper: Convert lists back to numpy arrays
        if 'joint_pos' in data:
            output_dict['joint_pos'] = np.array(data['joint_pos'], dtype=np.float64)
        
        if 'joint_names' in data:
            # joint_names must be numpy array of strings
            output_dict['joint_names'] = np.array(data['joint_names'])

        if 'base_pos_w' in data and data['base_pos_w']:
             output_dict['base_pos_w'] = np.array(data['base_pos_w'], dtype=np.float64)
             
        if 'base_quat_w' in data and data['base_quat_w']:
             output_dict['base_quat_w'] = np.array(data['base_quat_w'], dtype=np.float64)

        if 'framerate' in data:
            output_dict['framerate'] = np.array(data['framerate'], dtype=np.float64)

        # 3. Write to an in-memory buffer (like a virtual file)
        memory_file = io.BytesIO()
        np.savez(memory_file, **output_dict)
        memory_file.seek(0)

        # 4. Send the file back to the browser
        return send_file(
            memory_file,
            mimetype='application/octet-stream',
            as_attachment=True,
            download_name='edited_motion.npz'
        )

    except Exception as e:
        print(f"Error saving NPZ: {e}")
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run(debug=True, port=5000)