Update app.py

app.py

@@ -4,12 +4,12 @@
 import gradio as gr
 import cv2
 import numpy as np
-from PIL import Image
-import mediapipe as mp
 import os
 import json
 import pandas as pd
 from datetime import datetime
+from PIL import Image
+import mediapipe as mp
 
 # Setup folders
 os.makedirs("pose_images", exist_ok=True)
@@ -28,30 +28,11 @@ pose_model = mp_pose.Pose(static_image_mode=True, model_complexity=2)
 mp_drawing = mp.solutions.drawing_utils
 mp_styles = mp.solutions.drawing_styles
 
-# Define pose connections (edges) for graph structure
-POSE_CONNECTIONS = [
-    # Face
-    (0, 1), (1, 2), (2, 3), (3, 7), (0, 4), (4, 5), (5, 6), (6, 8),
-    # Torso
-    (9, 10), (11, 12), (11, 13), (13, 15), (15, 17), (15, 19), (15, 21),
-    (12, 14), (14, 16), (16, 18), (16, 20), (16, 22), (11, 23), (12, 24),
-    (23, 24),
-    # Left arm
-    (11, 13), (13, 15), (15, 17), (17, 19), (19, 21),
-    # Right arm
-    (12, 14), (14, 16), (16, 18), (18, 20), (20, 22),
-    # Left leg
-    (23, 25), (25, 27), (27, 29), (29, 31), (27, 31),
-    # Right leg
-    (24, 26), (26, 28), (28, 30), (30, 32), (28, 32)
-]
-
+# Define function to extract nodes and edges
 def create_pose_graph_data(pose_landmarks):
-    """Create nodes and edges data structure from pose landmarks"""
     nodes = {}
     edges = []
-    
-    # Create nodes
+
     for idx, lm in enumerate(pose_landmarks.landmark):
         name = mp_pose.PoseLandmark(idx).name
         nodes[idx] = {
@@ -62,11 +43,9 @@ def create_pose_graph_data(pose_landmarks):
             "z": round(lm.z, 4),
             "visibility": round(lm.visibility, 3)
         }
-    
-    # Create edges based on MediaPipe connections
+
     for connection in mp_pose.POSE_CONNECTIONS:
-        start_idx = connection[0]
-        end_idx = connection[1]
+        start_idx, end_idx = connection
         if start_idx < len(pose_landmarks.landmark) and end_idx < len(pose_landmarks.landmark):
             edges.append({
                 "from": start_idx,
@@ -74,27 +53,23 @@ def create_pose_graph_data(pose_landmarks):
                 "from_name": mp_pose.PoseLandmark(start_idx).name,
                 "to_name": mp_pose.PoseLandmark(end_idx).name
             })
-    
+
     return nodes, edges
 
+# Main pose processing function
 def process_pose(image, pose_description=""):
-    """Process pose image and return overlay with pose data"""
     if image is None:
         return None, "❌ Please upload an image.", "", None
-    
-    # Timestamp-based ID
+
     ts = datetime.now().strftime("%Y%m%d_%H%M%S")
     pose_id = f"pose_{ts}"
-    
-    # Convert to OpenCV
+
     img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-    
-    # Detect pose
     results = pose_model.process(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB))
+
     if not results.pose_landmarks:
         return None, "❌ No pose detected.", "", None
-    
-    # Draw overlay
+
     overlay = img_bgr.copy()
     mp_drawing.draw_landmarks(
         overlay,
@@ -102,18 +77,12 @@ def process_pose(image, pose_description=""):
         mp_pose.POSE_CONNECTIONS,
         landmark_drawing_spec=mp_styles.get_default_pose_landmarks_style()
     )
-    
-    # Convert back to RGB for display
     overlay_rgb = cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB)
-    
-    # Save overlay image
     overlay_path = f"pose_images/{pose_id}.png"
     cv2.imwrite(overlay_path, overlay)
-    
-    # Create graph data structure
+
     nodes, edges = create_pose_graph_data(results.pose_landmarks)
-    
-    # Create pose data summary
+
     pose_data = {
         "pose_id": pose_id,
         "total_nodes": len(nodes),
@@ -122,8 +91,7 @@ def process_pose(image, pose_description=""):
         "edges": edges,
         "description": pose_description if pose_description else "No description provided"
     }
-    
-    # Save to JSON dataset
+
     pose_dataset[pose_id] = {
         "pose_name": pose_id,
         "image_path": overlay_path,
@@ -131,242 +99,141 @@ def process_pose(image, pose_description=""):
         "pose_data": pose_data,
         "timestamp": ts
     }
-    
+
     with open(json_path, "w") as f:
         json.dump(pose_dataset, f, indent=2)
-    
-    # Format pose data for display
-    data_display = f"""
-🎯 **Pose Analysis Results**
+
+    data_display = f"""🎯 **Pose Analysis Results**
 
 📊 **Graph Structure:**
-- Total Nodes (Keypoints): {len(nodes)}
-- Total Edges (Connections): {len(edges)}
+- Total Nodes: {len(nodes)}
+- Total Edges: {len(edges)}
 
-📝 **Pose Description:** {pose_description if pose_description else "No description provided"}
+📝 **Pose Description:** {pose_description or "No description provided"}
 
 🔍 **Key Nodes (First 10):**
 """
-    
     for i, (idx, node) in enumerate(list(nodes.items())[:10]):
-        data_display += f"• {node['name']}: ({node['x']:.3f}, {node['y']:.3f}, {node['z']:.3f}) [visibility: {node['visibility']}]\n"
-    
+        data_display += f"• {node['name']}: ({node['x']}, {node['y']}, {node['z']}) [vis: {node['visibility']}]\n"
     if len(nodes) > 10:
         data_display += f"... and {len(nodes) - 10} more nodes\n"
-    
-    data_display += f"""
-🔗 **Sample Connections:**
-"""
-    
-    for i, edge in enumerate(edges[:5]):
+
+    data_display += "\n🔗 **Sample Edges:**\n"
+    for edge in edges[:5]:
         data_display += f"• {edge['from_name']} → {edge['to_name']}\n"
-    
     if len(edges) > 5:
         data_display += f"... and {len(edges) - 5} more connections\n"
-    
-    data_display += f"""
-💾 **Data Saved:**
-- Image: {overlay_path}
-- JSON: {json_path}
-- Pose ID: {pose_id}
-"""
-    
+
+    data_display += f"\n💾 **Saved as:** {overlay_path}"
+
     return overlay_rgb, data_display, f"✅ Pose '{pose_id}' saved successfully!", None
 
+# Save with description and return CSV
 def save_with_description(image, description):
-    """Save pose with description"""
     if image is None:
         return None, "❌ Please upload an image first.", "❌ No image to process", None
-    
+
     overlay_img, data_display, status, _ = process_pose(image, description)
     csv_file = create_csv_download()
     return overlay_img, data_display, status, csv_file
 
+# ✅ NEW Simplified CSV Creator (Only 3 fields)
 def create_csv_download():
-    """Create CSV file from pose dataset for download"""
     if not pose_dataset:
         return None
-    
-    # Prepare data for CSV
+
     csv_data = []
-    
     for pose_id, pose_info in pose_dataset.items():
-        pose_data = pose_info.get('pose_data', {})
-        nodes = pose_data.get('nodes', {})
-        edges = pose_data.get('edges', [])
-        
-        # Add nodes data
-        for node_id, node_info in nodes.items():
-            csv_data.append({
-                'pose_id': pose_id,
-                'pose_description': pose_info.get('pose_description', ''),
-                'timestamp': pose_info.get('timestamp', ''),
-                'data_type': 'node',
-                'element_id': node_id,
-                'element_name': node_info.get('name', ''),
-                'x_coordinate': node_info.get('x', ''),
-                'y_coordinate': node_info.get('y', ''),
-                'z_coordinate': node_info.get('z', ''),
-                'visibility': node_info.get('visibility', ''),
-                'connection_from': '',
-                'connection_to': '',
-                'connection_from_name': '',
-                'connection_to_name': ''
-            })
-        
-        # Add edges data
-        for edge in edges:
-            csv_data.append({
-                'pose_id': pose_id,
-                'pose_description': pose_info.get('pose_description', ''),
-                'timestamp': pose_info.get('timestamp', ''),
-                'data_type': 'edge',
-                'element_id': f"{edge['from']}-{edge['to']}",
-                'element_name': f"{edge['from_name']}_to_{edge['to_name']}",
-                'x_coordinate': '',
-                'y_coordinate': '',
-                'z_coordinate': '',
-                'visibility': '',
-                'connection_from': edge['from'],
-                'connection_to': edge['to'],
-                'connection_from_name': edge['from_name'],
-                'connection_to_name': edge['to_name']
-            })
-    
-    # Create DataFrame and save as CSV
+        pose_data = pose_info.get("pose_data", {})
+        pose_description = pose_info.get("pose_description", "")
+        image_path = pose_info.get("image_path", "")
+
+        csv_data.append({
+            "image_name": os.path.basename(image_path),
+            "pose_data": json.dumps(pose_data),
+            "pose_description": pose_description
+        })
+
     df = pd.DataFrame(csv_data)
-    csv_filename = f"pose_dataset_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
+    csv_filename = f"simplified_pose_dataset_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
     df.to_csv(csv_filename, index=False)
-    
+
     return csv_filename
 
+# Export only current pose
 def export_current_pose_csv(image, description=""):
-    """Export current pose analysis as CSV"""
     if image is None:
         return None, "❌ Please upload an image first."
-    
-    # Process current image
-    overlay_img, data_display, status = process_pose(image, description)
-    
-    # Create CSV for current pose only
+
+    overlay_img, data_display, status, _ = process_pose(image, description)
+
     if pose_dataset:
         latest_pose_id = max(pose_dataset.keys())
-        current_pose_data = {latest_pose_id: pose_dataset[latest_pose_id]}
-        
-        # Create CSV data for current pose
-        csv_data = []
-        pose_info = current_pose_data[latest_pose_id]
-        pose_data = pose_info.get('pose_data', {})
-        nodes = pose_data.get('nodes', {})
-        edges = pose_data.get('edges', [])
-        
-        # Add nodes data
-        for node_id, node_info in nodes.items():
-            csv_data.append({
-                'pose_id': latest_pose_id,
-                'pose_description': pose_info.get('pose_description', ''),
-                'timestamp': pose_info.get('timestamp', ''),
-                'data_type': 'node',
-                'element_id': node_id,
-                'element_name': node_info.get('name', ''),
-                'x_coordinate': node_info.get('x', ''),
-                'y_coordinate': node_info.get('y', ''),
-                'z_coordinate': node_info.get('z', ''),
-                'visibility': node_info.get('visibility', ''),
-                'connection_from': '',
-                'connection_to': '',
-                'connection_from_name': '',
-                'connection_to_name': ''
-            })
-        
-        # Add edges data
-        for edge in edges:
-            csv_data.append({
-                'pose_id': latest_pose_id,
-                'pose_description': pose_info.get('pose_description', ''),
-                'timestamp': pose_info.get('timestamp', ''),
-                'data_type': 'edge',
-                'element_id': f"{edge['from']}-{edge['to']}",
-                'element_name': f"{edge['from_name']}_to_{edge['to_name']}",
-                'x_coordinate': '',
-                'y_coordinate': '',
-                'z_coordinate': '',
-                'visibility': '',
-                'connection_from': edge['from'],
-                'connection_to': edge['to'],
-                'connection_from_name': edge['from_name'],
-                'connection_to_name': edge['to_name']
-            })
-        
-        # Create DataFrame and save as CSV
+        pose_info = pose_dataset[latest_pose_id]
+
+        csv_data = [{
+            "image_name": os.path.basename(pose_info.get("image_path", "")),
+            "pose_data": json.dumps(pose_info.get("pose_data", {})),
+            "pose_description": pose_info.get("pose_description", "")
+        }]
+
         df = pd.DataFrame(csv_data)
         csv_filename = f"current_pose_{latest_pose_id}.csv"
         df.to_csv(csv_filename, index=False)
-        
+
         return csv_filename, f"✅ Current pose exported as {csv_filename}"
-    
+
     return None, "❌ No pose data to export"
 
-# Create Gradio interface
-with gr.Blocks(title="🧘 Advanced Pose Analysis Tool") as demo:
-    gr.Markdown("# 🧘 Advanced Pose Analysis with MediaPipe")
-    gr.Markdown("Upload a yoga, archery, or any pose image to extract keypoints and analyze the pose structure as nodes and edges.")
-    
+# Gradio Interface
+with gr.Blocks(title="🧘 Simplified Pose Analysis Tool") as demo:
+    gr.Markdown("# 🧘 Pose Analysis with MediaPipe + CSV Export")
+    gr.Markdown("Upload a pose image and extract keypoints, description, and download results.")
+
     with gr.Row():
         with gr.Column(scale=1):
-            # Input section
-            gr.Markdown("## 📤 Input")
+            gr.Markdown("## 📤 Upload")
             input_image = gr.Image(type="numpy", label="Upload Pose Image")
             pose_description = gr.Textbox(
                 label="Pose Description",
-                placeholder="Enter a description of the pose (e.g., 'Warrior II pose with arms extended')",
+                placeholder="e.g. 'Triangle Pose with right arm up'",
                 lines=3
             )
-            
+
             with gr.Row():
                 analyze_btn = gr.Button("🔍 Analyze Pose", variant="primary")
                 save_btn = gr.Button("💾 Save with Description", variant="secondary")
-            
-            # CSV Download buttons
+
             gr.Markdown("### 📥 Download Options")
             with gr.Row():
                 download_current_btn = gr.Button("📄 Download Current Pose CSV", variant="secondary")
                 download_all_btn = gr.Button("📁 Download All Poses CSV", variant="secondary")
-            
-            # Download file outputs
+
             current_csv_download = gr.File(label="Current Pose CSV", visible=False)
             all_csv_download = gr.File(label="All Poses CSV", visible=False)
             download_status = gr.Textbox(label="Download Status", visible=False)
-        
+
         with gr.Column(scale=1):
-            # Output section
             gr.Markdown("## 📊 Results")
-            output_image = gr.Image(label="Pose with MediaPipe Overlay")
+            output_image = gr.Image(label="Pose with Overlay")
             status_text = gr.Textbox(label="Status", lines=1)
-    
-    # Pose data display
-    gr.Markdown("## 📋 Pose Data (Nodes & Edges)")
-    pose_data_display = gr.Textbox(
-        label="Pose Analysis Data",
-        lines=15,
-        max_lines=20,
-        show_copy_button=True
-    )
-    
-    # Button actions
+
+    gr.Markdown("## 🧠 Pose Data")
+    pose_data_display = gr.Textbox(label="Pose Details", lines=15, show_copy_button=True)
+
+    # Event handlers
     analyze_btn.click(
         fn=lambda img: process_pose(img, ""),
         inputs=[input_image],
         outputs=[output_image, pose_data_display, status_text, current_csv_download]
     )
-    
+
     save_btn.click(
         fn=save_with_description,
         inputs=[input_image, pose_description],
         outputs=[output_image, pose_data_display, status_text, current_csv_download]
     )
-    
-    # CSV Download actions
+
     download_current_btn.click(
         fn=export_current_pose_csv,
         inputs=[input_image, pose_description],
@@ -378,7 +245,7 @@ with gr.Blocks(title="🧘 Advanced Pose Analysis Tool") as demo:
         fn=lambda: gr.update(visible=True),
         outputs=[download_status]
     )
-    
+
     download_all_btn.click(
         fn=create_csv_download,
         outputs=[all_csv_download]
@@ -386,39 +253,25 @@ with gr.Blocks(title="🧘 Advanced Pose Analysis Tool") as demo:
         fn=lambda: gr.update(visible=True),
         outputs=[all_csv_download]
     )
-    
-    # Auto-analyze on image upload
+
     input_image.change(
         fn=lambda img: process_pose(img, ""),
         inputs=[input_image],
         outputs=[output_image, pose_data_display, status_text, current_csv_download]
     )
-    
+
     gr.Markdown("""
-    ## 📖 How to Use:
-    1. **Upload Image**: Upload a pose image using the image uploader
-    2. **Auto Analysis**: The pose will be automatically analyzed showing keypoints
-    3. **Add Description**: Enter a description of the pose in the text box
-    4. **Save**: Click "Save with Description" to save the pose data with your description
-    5. **Download CSV**: 
-       - "Download Current Pose CSV" - Export current pose data as CSV
-       - "Download All Poses CSV" - Export all saved poses as CSV
-    
-    ## 📊 Data Structure:
-    - **Nodes**: 33 body keypoints (face, torso, arms, legs) with 3D coordinates
-    - **Edges**: Connections between keypoints following human body structure
-    - **Visibility**: Confidence score for each keypoint detection
-    
-    ## 📁 CSV Format:
-    The CSV file contains columns:
-    - `pose_id`: Unique identifier for each pose
-    - `pose_description`: Your description of the pose
-    - `data_type`: 'node' for keypoints, 'edge' for connections
-    - `element_name`: Name of the body part or connection
-    - `x_coordinate`, `y_coordinate`, `z_coordinate`: 3D position
-    - `visibility`: Detection confidence (for nodes)
-    - `connection_from/to`: Start/end points (for edges)
-    """)
-
-# Launch the interface
-demo.launch(share=True)
+## 📝 How It Works
+1. Upload a pose image
+2. Automatically analyze and overlay pose keypoints
+3. Add an optional description
+4. Save and export to CSV
+
+### CSV Format (Simplified):
+- `image_name`: Name of the saved pose image
+- `pose_data`: JSON string containing nodes and edges
+- `pose_description`: Your text description
+""")
+
+# Launch app
+demo.launch(share=True)