app.py

import streamlit as st
import subprocess
import sys
import os
import time
import shutil
import tempfile
import threading
import queue
from datetime import datetime
from pathlib import Path

def run_command(command, working_dir, progress_bar, progress_text, step_start_progress, step_weight, show_progress=True):
    try:
        env = os.environ.copy()
        env["PYTHONUNBUFFERED"] = "1"
        
        process = subprocess.Popen(
            command,
            cwd=working_dir,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1,
            universal_newlines=True,
            env=env
        )
        
        stdout_queue = queue.Queue()
        stderr_queue = queue.Queue()
        
        def read_output(pipe, q, source):
            for line in iter(pipe.readline, ''):
                q.put((source, line))
            pipe.close()
        
        stdout_thread = threading.Thread(target=read_output, args=(process.stdout, stdout_queue, 'stdout'))
        stderr_thread = threading.Thread(target=read_output, args=(process.stderr, stderr_queue, 'stderr'))
        stdout_thread.daemon = True
        stderr_thread.daemon = True
        stdout_thread.start()
        stderr_thread.start()
        
        total_progress = step_start_progress
        stderr_lines = []
        
        while process.poll() is None or not (stdout_queue.empty() and stderr_queue.empty()):
            try:
                source, line = next((q.get_nowait() for q in [stdout_queue, stderr_queue] if not q.empty()), (None, None))
                if line:
                    if source == 'stdout':
                        if show_progress and line.startswith("PROGRESS:"):
                            try:
                                progress = float(line.strip().split("PROGRESS:")[1]) / 100
                                if Path(command[1]).name == 'gen_skes.py':
                                    # Scale 0-200% to 0-80%
                                    adjusted_progress = step_start_progress + (progress / 2) * step_weight
                                else:
                                    # Scale 0-100% to 20% for conver_bvh.py
                                    adjusted_progress = step_start_progress + (progress * step_weight)
                                total_progress = min(adjusted_progress, step_start_progress + step_weight)
                                progress_bar.progress(total_progress)
                                progress_text.text(f"{int(total_progress * 100)}%")
                            except ValueError:
                                pass
                    elif source == 'stderr':
                        stderr_lines.append(line.strip())
            except queue.Empty:
                time.sleep(0.01)
        
        stdout_thread.join()
        stderr_thread.join()
        
        if process.returncode != 0:
            stderr_output = '\n'.join(stderr_lines)
            st.error(f"Error in {Path(command[1]).name}:\n{stderr_output}")
            return False
        
        if show_progress:
            progress_bar.progress(step_start_progress + step_weight)
            progress_text.text(f"{int((step_start_progress + step_weight) * 100)}%")
        return True
    
    except Exception as e:
        st.error(f"Exception in {Path(command[1]).name}: {str(e)}")
        return False

def cleanup_output_folder(output_dir, delay=1800):
    time.sleep(delay)
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir, ignore_errors=True)
        print(f"Deleted temporary output folder after timeout: {output_dir}")

def process_video(video_file):
    base_dir = Path(__file__).parent.resolve()
    
    gen_skes_path = base_dir / "VideoToNPZ" / "gen_skes.py"
    convert_obj_path = base_dir / "convertNPZtoBVH" / "conver_obj.py"
    convert_bvh_path = base_dir / "convertNPZtoBVH" / "conver_bvh.py"
    
    for script_path in [gen_skes_path, convert_obj_path, convert_bvh_path]:
        if not script_path.exists():
            st.error(f"Required script not found: {script_path}")
            return None
    
    with tempfile.TemporaryDirectory() as tmp_dir:
        video_path = Path(tmp_dir) / "input_video.mp4"
        with open(video_path, "wb") as f:
            f.write(video_file.read())
        
        if not video_path.exists():
            st.error(f"Video file not found at: {video_path}")
            return None
        
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        output_dir = base_dir / f"outputs_{timestamp}"
        output_dir.mkdir(exist_ok=True)
        
        if not os.access(output_dir, os.W_OK):
            st.error(f"Cannot write to output directory: {output_dir}")
            return None
        
        default_output_dir = base_dir / "outputs"
        
        pipeline_steps = [
            {"command": [sys.executable, str(gen_skes_path), "-v", str(video_path)], "working_dir": gen_skes_path.parent, "weight": 0.8, "show_progress": True},
            {"command": [sys.executable, str(convert_obj_path), "--output-dir", str(output_dir)], "working_dir": convert_obj_path.parent, "weight": 0.0, "show_progress": False},
            {"command": [sys.executable, str(convert_bvh_path), "--output-dir", str(output_dir)], "working_dir": convert_bvh_path.parent, "weight": 0.2, "show_progress": True}
        ]
        
        progress_bar = st.progress(0.0)
        progress_text = st.empty()  # Added for percentage display
        total_progress = 0.0
        
        for i, step in enumerate(pipeline_steps):
            success = run_command(
                step["command"], 
                step["working_dir"], 
                progress_bar, 
                progress_text, 
                total_progress, 
                step["weight"], 
                show_progress=step["show_progress"]
            )
            if not success:
                st.error(f"Failed at step: {' '.join(map(str, step['command']))}")
                if default_output_dir.exists():
                    shutil.rmtree(default_output_dir, ignore_errors=True)
                return None
            
            if i == 0 and default_output_dir.exists():
                npz_dir = default_output_dir / "npz"
                if npz_dir.exists():
                    target_npz_dir = output_dir / "npz"
                    shutil.move(str(npz_dir), str(target_npz_dir))
                if default_output_dir.exists():
                    shutil.rmtree(default_output_dir, ignore_errors=True)
            
            total_progress += step["weight"]
            if step["show_progress"]:
                progress_bar.progress(min(total_progress, 1.0))
                progress_text.text(f"{int(total_progress * 100)}%")
        
        bvh_output_dir = output_dir / "bvh"
        bvh_file = bvh_output_dir / "output.bvh"
        
        if bvh_file.exists():
            cleanup_thread = threading.Thread(target=cleanup_output_folder, args=(output_dir,))
            cleanup_thread.daemon = True
            cleanup_thread.start()
            
            return {
                'bvh_file': bvh_file,
                'output_dir': output_dir
            }
        else:
            st.error(f"Failed to generate BVH file at: {bvh_file}")
            if default_output_dir.exists():
                shutil.rmtree(default_output_dir, ignore_errors=True)
            return None

def cleanup_immediate(output_dir):
    if output_dir and os.path.exists(output_dir):
        shutil.rmtree(output_dir, ignore_errors=True)
        st.success("")
    else:
        st.warning("")

def main():
    st.title("Video to BVH Converter")
    st.write("Upload a video file to convert it to BVH format")
    
    uploaded_file = st.file_uploader("Choose a video file", type=['mp4', 'avi', 'mov'])
    
    if uploaded_file is not None:
        st.video(uploaded_file)
        
        if st.button("Convert to BVH"):
            with st.spinner("Processing video..."):
                conversion_result = process_video(uploaded_file)
                
                if conversion_result:
                    with open(conversion_result['bvh_file'], "rb") as f:
                        st.download_button(
                            label="Download BVH File",
                            data=f,
                            file_name="output.bvh",
                            mime="application/octet-stream",
                            on_click=cleanup_immediate,
                            args=(conversion_result['output_dir'],)
                        )
                    st.success("Conversion completed successfully! File will be deleted after download or in 30 minutes.")
    
    st.markdown("""
    ### Instructions
    1. Upload a video file using the uploader above
    2. Click 'Convert to BVH' to start the conversion process
    3. Watch the progress bar as it processes
    4. Download the resulting BVH file - folder will be deleted immediately after download or after 30 minutes
    """)

if __name__ == "__main__":
    main()