Update app.py

app.py

@@ -91,11 +91,11 @@ def get_bvh_from_folder(drive_service, folder_id=None):
             bvh_file = files[0]
             bvh_id = bvh_file['id']
             bvh_url = f"https://drive.google.com/uc?id={bvh_id}"
-            return bvh_id, bvh_url
-        return None, None
+            return bvh_id, bvh_url, bvh_file['name']
+        return None, None, None
     except Exception as e:
         st.error(f"Error checking folder for BVH: {e}")
-        return None, None
+        return None, None, None
 
 def download_notebook_from_drive(drive_service, temp_dir):
     try:
@@ -197,6 +197,22 @@ def check_kernel_exists(api, notebook_slug):
         st.error(f"Kernel check failed: {e}")
         return False
 
+def download_and_save_bvh(bvh_url, filename):
+    try:
+        response = requests.get(bvh_url)
+        if response.status_code == 200:
+            temp_dir = tempfile.mkdtemp()
+            bvh_path = os.path.join(temp_dir, filename)
+            with open(bvh_path, 'wb') as f:
+                f.write(response.content)
+            return bvh_path, response.content
+        else:
+            st.error(f"Failed to download BVH: Status code {response.status_code}")
+            return None, None
+    except Exception as e:
+        st.error(f"Error downloading BVH: {e}")
+        return None, None
+
 def process_video(api, drive_service, video_file):
     video_file_id = None
     bvh_file_id = None
@@ -259,18 +275,32 @@ def process_video(api, drive_service, video_file):
                                 delete_from_drive(drive_service, video_file_id)
                             return None
 
-                    bvh_file_id, bvh_url = get_bvh_from_folder(drive_service)
-                    if bvh_url:
-                        response = requests.get(bvh_url)
-                        if response.status_code == 200:
+                    bvh_file_id, bvh_url, bvh_filename = get_bvh_from_folder(drive_service)
+                    if bvh_url and bvh_filename:
+                        # Automatically download the BVH file
+                        bvh_path, bvh_data = download_and_save_bvh(bvh_url, bvh_filename or "motion_capture.bvh")
+                        
+                        if bvh_path and bvh_data:
                             progress_bar.progress(1.0)
                             progress_text.text("Complete!")
+                            
+                            # Cleanup files on Google Drive
                             if video_file_id:
                                 delete_from_drive(drive_service, video_file_id)
                             if bvh_file_id:
                                 delete_from_drive(drive_service, bvh_file_id)
-                            st.write("BVH generated")
-                            return {'bvh_data': response.content}
+                            
+                            st.success("Motion capture complete! BVH file ready for download.")
+                            
+                            # Save the BVH data to session state for download
+                            st.session_state['bvh_data'] = bvh_data
+                            st.session_state['bvh_filename'] = bvh_filename or "motion_capture.bvh"
+                            
+                            return {
+                                'bvh_data': bvh_data,
+                                'bvh_path': bvh_path,
+                                'bvh_filename': bvh_filename or "motion_capture.bvh"
+                            }
 
                     if execution_started and current_status in ['complete', 'error']:
                         progress_bar.progress(0.8 if current_status == 'complete' else 0.6)
@@ -381,15 +411,26 @@ def main():
             result = process_video(api, drive_service, st.session_state['uploaded_file'])
             
             if result and 'bvh_data' in result:
-                st.success("Motion capture complete!")
+                # The BVH file has already been downloaded and deleted from Drive
+                # Now just offer it for download to the user
                 st.download_button(
                     label="Download BVH",
                     data=result['bvh_data'],
-                    file_name="motion_capture.bvh",
+                    file_name=result['bvh_filename'],
                     mime="application/octet-stream"
                 )
             else:
-                st.error("Failed to retrieve BVH file.")
+                st.error("Failed to generate BVH file.")
+    
+    # If BVH data is in session state (from a previous run), offer it for download
+    if 'bvh_data' in st.session_state and 'bvh_filename' in st.session_state:
+        st.download_button(
+            label="Download BVH",
+            data=st.session_state['bvh_data'],
+            file_name=st.session_state['bvh_filename'],
+            mime="application/octet-stream"
+        )
+                
     st.markdown('</div>', unsafe_allow_html=True)
 
 if __name__ == "__main__":