app.py

import gradio as gr
import json
import requests
import os
import math
from PIL import Image, ImageDraw, ImageFont
import numpy as np

# Fireworks AI configuration
FIREWORKS_API_KEY = os.getenv("FIREWORKS_API_KEY", "YOUR_API_KEY_HERE")
FIREWORKS_API_URL = "https://api.fireworks.ai/inference/v1/chat/completions"

# OpenPose keypoint definitions
BODY_PARTS = {
    "Nose": 0, "Neck": 1, "RShoulder": 2, "RElbow": 3, "RWrist": 4,
    "LShoulder": 5, "LElbow": 6, "LWrist": 7, "RHip": 8, "RKnee": 9,
    "RAnkle": 10, "LHip": 11, "LKnee": 12, "LAnkle": 13, "REye": 14,
    "LEye": 15, "REar": 16, "LEar": 17
}

# Enhanced skeleton connections with thickness
POSE_CONNECTIONS = [
    ("Neck", "Nose", 2), ("Nose", "REye", 1), ("REye", "REar", 1), 
    ("Nose", "LEye", 1), ("LEye", "LEar", 1),
    ("Neck", "RShoulder", 3), ("RShoulder", "RElbow", 3), ("RElbow", "RWrist", 2),
    ("Neck", "LShoulder", 3), ("LShoulder", "LElbow", 3), ("LElbow", "LWrist", 2),
    ("Neck", "RHip", 4), ("RHip", "RKnee", 4), ("RKnee", "RAnkle", 3),
    ("Neck", "LHip", 4), ("LHip", "LKnee", 4), ("LKnee", "LAnkle", 3),
    ("RHip", "LHip", 4)
]

# Dynamic and precise pose templates
DYNAMIC_POSES = {
    "🏃 Sprint": {
        "Nose": [256, 60], "Neck": [256, 100],
        "RShoulder": [230, 120], "RElbow": [260, 140], "RWrist": [290, 120],
        "LShoulder": [282, 120], "LElbow": [252, 160], "LWrist": [222, 200],
        "RHip": [245, 260], "RKnee": [220, 360], "RAnkle": [195, 460],
        "LHip": [267, 260], "LKnee": [300, 340], "LAnkle": [330, 420],
        "REye": [246, 50], "LEye": [266, 50], "REar": [236, 55], "LEar": [276, 55]
    },
    "🤸 Backflip": {
        "Nose": [256, 250], "Neck": [256, 280],
        "RShoulder": [220, 300], "RElbow": [180, 280], "RWrist": [140, 260],
        "LShoulder": [292, 300], "LElbow": [332, 280], "LWrist": [372, 260],
        "RHip": [240, 180], "RKnee": [220, 120], "RAnkle": [200, 60],
        "LHip": [272, 180], "LKnee": [292, 120], "LAnkle": [312, 60],
        "REye": [246, 260], "LEye": [266, 260], "REar": [236, 255], "LEar": [276, 255]
    },
    "🥋 Martial Arts Kick": {
        "Nose": [256, 80], "Neck": [256, 120],
        "RShoulder": [220, 140], "RElbow": [180, 120], "RWrist": [140, 100],
        "LShoulder": [292, 140], "LElbow": [332, 160], "LWrist": [372, 180],
        "RHip": [240, 280], "RKnee": [235, 380], "RAnkle": [230, 480],
        "LHip": [272, 280], "LKnee": [350, 250], "LAnkle": [420, 220],
        "REye": [246, 70], "LEye": [266, 70], "REar": [236, 75], "LEar": [276, 75]
    },
    "🩰 Ballet Leap": {
        "Nose": [256, 100], "Neck": [256, 140],
        "RShoulder": [200, 130], "RElbow": [150, 110], "RWrist": [100, 90],
        "LShoulder": [312, 130], "LElbow": [362, 110], "LWrist": [412, 90],
        "RHip": [240, 300], "RKnee": [180, 380], "RAnkle": [120, 460],
        "LHip": [272, 300], "LKnee": [332, 380], "LAnkle": [392, 460],
        "REye": [246, 90], "LEye": [266, 90], "REar": [236, 95], "LEar": [276, 95]
    },
    "🏀 Basketball Dunk": {
        "Nose": [256, 60], "Neck": [256, 100],
        "RShoulder": [220, 120], "RElbow": [200, 80], "RWrist": [180, 40],
        "LShoulder": [292, 120], "LElbow": [312, 80], "LWrist": [332, 40],
        "RHip": [240, 260], "RKnee": [230, 360], "RAnkle": [225, 460],
        "LHip": [272, 260], "LKnee": [282, 340], "LAnkle": [287, 420],
        "REye": [246, 50], "LEye": [266, 50], "REar": [236, 55], "LEar": [276, 55]
    },
    "🧘 Advanced Yoga": {
        "Nose": [256, 380], "Neck": [256, 420],
        "RShoulder": [220, 440], "RElbow": [200, 400], "RWrist": [180, 360],
        "LShoulder": [292, 440], "LElbow": [312, 400], "LWrist": [332, 360],
        "RHip": [240, 280], "RKnee": [220, 180], "RAnkle": [200, 80],
        "LHip": [272, 280], "LKnee": [292, 180], "LAnkle": [312, 80],
        "REye": [246, 390], "LEye": [266, 390], "REar": [236, 385], "LEar": [276, 385]
    },
    "🕺 Breakdance Freeze": {
        "Nose": [150, 400], "Neck": [180, 420],
        "RShoulder": [200, 400], "RElbow": [220, 350], "RWrist": [240, 300],
        "LShoulder": [160, 440], "LElbow": [140, 480], "LWrist": [120, 500],
        "RHip": [300, 350], "RKnee": [350, 300], "RAnkle": [400, 250],
        "LHip": [280, 380], "LKnee": [330, 420], "LAnkle": [380, 460],
        "REye": [140, 390], "LEye": [160, 390], "REar": [130, 395], "LEar": [170, 395]
    },
    "🏋️ Weightlifting": {
        "Nose": [256, 100], "Neck": [256, 140],
        "RShoulder": [210, 160], "RElbow": [180, 120], "RWrist": [150, 80],
        "LShoulder": [302, 160], "LElbow": [332, 120], "LWrist": [362, 80],
        "RHip": [235, 300], "RKnee": [220, 400], "RAnkle": [210, 480],
        "LHip": [277, 300], "LKnee": [292, 400], "LAnkle": [302, 480],
        "REye": [246, 90], "LEye": [266, 90], "REar": [236, 95], "LEar": [276, 95]
    },
    "🏊 Swimming Dive": {
        "Nose": [380, 200], "Neck": [360, 220],
        "RShoulder": [340, 240], "RElbow": [380, 230], "RWrist": [420, 220],
        "LShoulder": [340, 200], "LElbow": [380, 190], "LWrist": [420, 180],
        "RHip": [280, 260], "RKnee": [220, 280], "RAnkle": [160, 300],
        "LHip": [280, 240], "LKnee": [220, 250], "LAnkle": [160, 260],
        "REye": [390, 195], "LEye": [390, 205], "REar": [385, 190], "LEar": [385, 210]
    },
    "🎾 Tennis Serve": {
        "Nose": [256, 80], "Neck": [256, 120],
        "RShoulder": [220, 140], "RElbow": [200, 100], "RWrist": [180, 60],
        "LShoulder": [292, 140], "LElbow": [312, 180], "LWrist": [332, 220],
        "RHip": [240, 280], "RKnee": [235, 380], "RAnkle": [230, 480],
        "LHip": [272, 280], "LKnee": [277, 380], "LAnkle": [282, 480],
        "REye": [246, 70], "LEye": [266, 70], "REar": [236, 75], "LEar": [276, 75]
    },
    "🪂 Skydiving": {
        "Nose": [256, 200], "Neck": [256, 240],
        "RShoulder": [200, 260], "RElbow": [150, 240], "RWrist": [100, 220],
        "LShoulder": [312, 260], "LElbow": [362, 240], "LWrist": [412, 220],
        "RHip": [240, 340], "RKnee": [200, 400], "RAnkle": [160, 460],
        "LHip": [272, 340], "LKnee": [312, 400], "LAnkle": [352, 460],
        "REye": [246, 190], "LEye": [266, 190], "REar": [236, 195], "LEar": [276, 195]
    },
    "🤾 Handball Jump": {
        "Nose": [256, 120], "Neck": [256, 160],
        "RShoulder": [220, 180], "RElbow": [200, 140], "RWrist": [180, 100],
        "LShoulder": [292, 180], "LElbow": [312, 220], "LWrist": [332, 260],
        "RHip": [240, 320], "RKnee": [220, 400], "RAnkle": [200, 480],
        "LHip": [272, 320], "LKnee": [292, 380], "LAnkle": [312, 440],
        "REye": [246, 110], "LEye": [266, 110], "REar": [236, 115], "LEar": [276, 115]
    }
}

def rotate_point(point, center, angle):
    """Rotate a point around a center by angle (in radians)"""
    x, y = point
    cx, cy = center
    cos_a = math.cos(angle)
    sin_a = math.sin(angle)
    nx = cos_a * (x - cx) - sin_a * (y - cy) + cx
    ny = sin_a * (x - cx) + cos_a * (y - cy) + cy
    return [nx, ny]

def scale_pose(keypoints, scale_factor):
    """Scale pose around center"""
    if not keypoints:
        return keypoints
    
    # Find center
    valid_points = [p for p in keypoints.values() if p]
    if not valid_points:
        return keypoints
    
    cx = sum(p[0] for p in valid_points) / len(valid_points)
    cy = sum(p[1] for p in valid_points) / len(valid_points)
    
    # Scale around center
    scaled = {}
    for part, point in keypoints.items():
        if point:
            x, y = point
            nx = cx + (x - cx) * scale_factor
            ny = cy + (y - cy) * scale_factor
            scaled[part] = [nx, ny]
        else:
            scaled[part] = point
    
    return scaled

def add_motion_blur(keypoints, direction="horizontal", intensity=0.3):
    """Add motion effect to pose"""
    blurred = keypoints.copy()
    
    for part in ["RWrist", "LWrist", "RAnkle", "LAnkle"]:
        if part in blurred and blurred[part]:
            x, y = blurred[part]
            if direction == "horizontal":
                blurred[part] = [x + intensity * 30, y]
            elif direction == "vertical":
                blurred[part] = [x, y + intensity * 30]
    
    return blurred

def draw_pose(keypoints, width=512, height=512, style="dynamic"):
    """Draw pose with enhanced visualization"""
    img = Image.new('RGB', (width, height), color='white')
    draw = ImageDraw.Draw(img)
    
    # Dynamic gradient background for action poses
    if style == "dynamic":
        for i in range(height):
            color_val = 255 - int((i / height) * 50)
            draw.rectangle([(0, i), (width, i+1)], fill=(color_val, color_val, 255))
    
    # Draw shadow/trail effect for dynamic poses
    if style == "dynamic":
        for start, end, _ in POSE_CONNECTIONS:
            if start in keypoints and end in keypoints:
                start_point = keypoints[start]
                end_point = keypoints[end]
                if start_point and end_point:
                    # Shadow effect
                    shadow_start = [start_point[0] + 5, start_point[1] + 5]
                    shadow_end = [end_point[0] + 5, end_point[1] + 5]
                    draw.line([tuple(shadow_start), tuple(shadow_end)], 
                            fill=(200, 200, 200), width=2)
    
    # Draw skeleton connections with variable thickness
    for connection in POSE_CONNECTIONS:
        start, end = connection[0], connection[1]
        thickness = connection[2] if len(connection) > 2 else 3
        
        if start in keypoints and end in keypoints:
            start_point = keypoints[start]
            end_point = keypoints[end]
            if start_point and end_point:
                # Main line
                draw.line([tuple(start_point), tuple(end_point)], 
                         fill='darkblue', width=thickness)
                # Highlight
                draw.line([tuple(start_point), tuple(end_point)], 
                         fill='blue', width=thickness-1)
    
    # Draw keypoints with gradient effect
    for part, point in keypoints.items():
        if point:
            x, y = point
            # Determine size based on body part
            if part in ["Neck", "RHip", "LHip"]:
                radius = 7
            elif part in ["RWrist", "LWrist", "RAnkle", "LAnkle"]:
                radius = 4
            else:
                radius = 5
            
            # Outer circle
            draw.ellipse([x-radius-1, y-radius-1, x+radius+1, y+radius+1], 
                        fill='darkred', outline='black')
            # Inner circle
            draw.ellipse([x-radius, y-radius, x+radius, y+radius], 
                        fill='red', outline='darkred')
            # Highlight
            draw.ellipse([x-radius+2, y-radius+2, x+radius-2, y+radius-2], 
                        fill='pink', outline=None)
    
    return img

def generate_pose_from_llm(prompt):
    """Generate dynamic pose using LLM with enhanced prompt"""
    enhanced_prompt = f"""You are an expert in generating DYNAMIC and PRECISE human pose keypoints.
    
    Task: Generate highly dynamic, action-oriented pose for: {prompt}
    
    Requirements:
    1. Canvas: 512x512 pixels
    2. Create EXTREME and DYNAMIC poses with:
       - Large range of motion
       - Athletic and acrobatic positions
       - Dramatic angles and perspectives
       - Action-oriented body positions
    3. Use all 18 keypoints: Nose, Neck, RShoulder, RElbow, RWrist, LShoulder, LElbow, LWrist,
       RHip, RKnee, RAnkle, LHip, LKnee, LAnkle, REye, LEye, REar, LEar
    4. Make poses that show:
       - Movement and energy
       - Athletic performance
       - Extreme flexibility
       - Dramatic action
    
    Return ONLY a JSON object with keypoint names and [x, y] coordinates.
    Make it VERY dynamic and exciting!"""
    
    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": f"Bearer {FIREWORKS_API_KEY}"
    }
    
    payload = {
        "model": "accounts/fireworks/models/qwen3-235b-a22b-instruct-2507",
        "max_tokens": 1024,
        "temperature": 0.7,  # Higher temperature for more creative poses
        "messages": [
            {"role": "system", "content": enhanced_prompt},
            {"role": "user", "content": f"Create an extremely dynamic pose for: {prompt}"}
        ]
    }
    
    try:
        response = requests.post(FIREWORKS_API_URL, headers=headers, json=payload, timeout=30)
        if response.status_code == 200:
            data = response.json()
            content = data['choices'][0]['message']['content']
            
            import re
            json_match = re.search(r'\{.*\}', content, re.DOTALL)
            if json_match:
                keypoints = json.loads(json_match.group())
                return keypoints
    except Exception as e:
        print(f"LLM Error: {e}")
    
    return get_dynamic_template(prompt)

def get_dynamic_template(prompt):
    """Select dynamic template based on prompt"""
    if not prompt:
        return DYNAMIC_POSES["🏃 Sprint"]
    
    prompt_lower = prompt.lower()
    
    # Check for specific action keywords
    if any(word in prompt_lower for word in ["run", "sprint", "dash"]):
        return DYNAMIC_POSES["🏃 Sprint"]
    elif any(word in prompt_lower for word in ["flip", "backflip", "acrobat"]):
        return DYNAMIC_POSES["🤸 Backflip"]
    elif any(word in prompt_lower for word in ["kick", "martial", "karate", "taekwondo"]):
        return DYNAMIC_POSES["🥋 Martial Arts Kick"]
    elif any(word in prompt_lower for word in ["ballet", "dance", "leap", "jump"]):
        return DYNAMIC_POSES["🩰 Ballet Leap"]
    elif any(word in prompt_lower for word in ["basketball", "dunk", "shoot"]):
        return DYNAMIC_POSES["🏀 Basketball Dunk"]
    elif any(word in prompt_lower for word in ["yoga", "stretch", "flexible"]):
        return DYNAMIC_POSES["🧘 Advanced Yoga"]
    elif any(word in prompt_lower for word in ["breakdance", "freeze", "bboy"]):
        return DYNAMIC_POSES["🕺 Breakdance Freeze"]
    elif any(word in prompt_lower for word in ["weight", "lift", "gym"]):
        return DYNAMIC_POSES["🏋️ Weightlifting"]
    elif any(word in prompt_lower for word in ["swim", "dive", "water"]):
        return DYNAMIC_POSES["🏊 Swimming Dive"]
    elif any(word in prompt_lower for word in ["tennis", "serve", "racket"]):
        return DYNAMIC_POSES["🎾 Tennis Serve"]
    elif any(word in prompt_lower for word in ["sky", "fall", "fly"]):
        return DYNAMIC_POSES["🪂 Skydiving"]
    elif any(word in prompt_lower for word in ["handball", "throw", "ball"]):
        return DYNAMIC_POSES["🤾 Handball Jump"]
    else:
        # Return random dynamic pose
        import random
        return random.choice(list(DYNAMIC_POSES.values()))

def apply_physics(keypoints, effect="gravity"):
    """Apply physics effects to make poses more realistic"""
    modified = keypoints.copy()
    
    if effect == "gravity":
        # Apply downward pull to limbs
        for part in ["RWrist", "LWrist"]:
            if part in modified and modified[part]:
                modified[part][1] += 10
    elif effect == "momentum":
        # Add motion direction
        if "RWrist" in modified and modified["RWrist"]:
            modified["RWrist"][0] += 15
        if "LAnkle" in modified and modified["LAnkle"]:
            modified["LAnkle"][0] -= 10
    
    return modified

def refine_pose_advanced(current_keypoints, instruction):
    """Advanced pose refinement with multiple options"""
    if not current_keypoints or not instruction:
        return current_keypoints
    
    keypoints = current_keypoints.copy()
    instruction_lower = instruction.lower()
    
    # Complex refinements
    if "rotate" in instruction_lower:
        angle = math.pi / 6  # 30 degrees
        center = [256, 256]
        for part, point in keypoints.items():
            if point:
                keypoints[part] = rotate_point(point, center, angle)
    
    elif "scale" in instruction_lower:
        if "up" in instruction_lower or "large" in instruction_lower:
            keypoints = scale_pose(keypoints, 1.2)
        elif "down" in instruction_lower or "small" in instruction_lower:
            keypoints = scale_pose(keypoints, 0.8)
    
    elif "jump" in instruction_lower or "leap" in instruction_lower:
        # Move everything up
        for part, point in keypoints.items():
            if point:
                keypoints[part][1] -= 50
        # Spread legs
        if "RAnkle" in keypoints:
            keypoints["RAnkle"][0] -= 30
        if "LAnkle" in keypoints:
            keypoints["LAnkle"][0] += 30
    
    elif "crouch" in instruction_lower or "squat" in instruction_lower:
        # Lower body parts
        for part in ["RKnee", "LKnee", "RAnkle", "LAnkle"]:
            if part in keypoints and keypoints[part]:
                keypoints[part][1] += 40
    
    elif "spread" in instruction_lower:
        if "arm" in instruction_lower:
            if "RWrist" in keypoints:
                keypoints["RWrist"][0] -= 40
            if "LWrist" in keypoints:
                keypoints["LWrist"][0] += 40
        elif "leg" in instruction_lower:
            if "RAnkle" in keypoints:
                keypoints["RAnkle"][0] -= 40
            if "LAnkle" in keypoints:
                keypoints["LAnkle"][0] += 40
    
    elif "twist" in instruction_lower:
        # Create twisting effect
        if "RShoulder" in keypoints:
            keypoints["RShoulder"][0] += 20
        if "LShoulder" in keypoints:
            keypoints["LShoulder"][0] -= 20
        if "RHip" in keypoints:
            keypoints["RHip"][0] -= 15
        if "LHip" in keypoints:
            keypoints["LHip"][0] += 15
    
    return keypoints

def keypoints_to_openpose_format(keypoints):
    """Convert to OpenPose JSON format with confidence scores"""
    if not keypoints:
        return "{}"
    
    candidate = []
    for i in range(18):
        part_name = None
        for name, idx in BODY_PARTS.items():
            if idx == i:
                part_name = name
                break
        
        if part_name and part_name in keypoints:
            x, y = keypoints[part_name]
            candidate.append([float(x), float(y), 1.0])
        else:
            candidate.append([0.0, 0.0, 0.0])
    
    subset = [[i for i in range(18) if candidate[i][2] > 0]]
    subset[0].extend([float(len(subset[0])), len(subset[0])])
    
    return json.dumps({"candidate": candidate, "subset": subset}, indent=2)

# Main generation function
def generate_pose(prompt, use_llm, template, physics_effect):
    """Generate dynamic pose with physics"""
    if template and template != "None":
        keypoints = DYNAMIC_POSES[template]
    elif use_llm and FIREWORKS_API_KEY != "YOUR_API_KEY_HERE" and prompt:
        keypoints = generate_pose_from_llm(prompt)
    elif prompt:
        keypoints = get_dynamic_template(prompt)
    else:
        import random
        keypoints = random.choice(list(DYNAMIC_POSES.values()))
    
    # Apply physics effect
    if physics_effect != "None":
        keypoints = apply_physics(keypoints, physics_effect.lower())
    
    pose_img = draw_pose(keypoints, style="dynamic")
    json_str = keypoints_to_openpose_format(keypoints)
    
    return pose_img, json_str, keypoints

def refine_existing_pose(instruction, keypoints_json):
    """Refine pose with advanced options"""
    if not keypoints_json:
        return None, "{}", {}
    
    refined_keypoints = refine_pose_advanced(keypoints_json, instruction)
    pose_img = draw_pose(refined_keypoints, style="dynamic")
    json_str = keypoints_to_openpose_format(refined_keypoints)
    
    return pose_img, json_str, refined_keypoints

def check_api_status():
    """Check API status"""
    if FIREWORKS_API_KEY != "YOUR_API_KEY_HERE":
        return "✅ API key configured - Advanced AI ready"
    return "⚠️ API key not set - Using dynamic templates"

# Create Gradio interface
app = gr.Blocks(title="Dynamic AI Pose Generator", theme=gr.themes.Soft())

with app:
    keypoints_state = gr.State({})
    
    gr.Markdown("""
    # 🎯 Dynamic AI Pose Generator
    ### Generate extremely dynamic and precise action poses!
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            api_status = gr.Markdown(check_api_status())
            
            use_llm = gr.Checkbox(
                label="🚀 Use Advanced AI (Fireworks)",
                value=False
            )
            
            prompt = gr.Textbox(
                label="Describe your action pose",
                placeholder="e.g., Ninja doing a flying kick, Athlete jumping over hurdles",
                lines=2
            )
            
            with gr.Row():
                template = gr.Dropdown(
                    choices=["None"] + list(DYNAMIC_POSES.keys()),
                    label="🎬 Dynamic Templates",
                    value="None"
                )
                
                physics_effect = gr.Dropdown(
                    choices=["None", "Gravity", "Momentum"],
                    label="⚡ Physics Effect",
                    value="None"
                )
            
            generate_btn = gr.Button("💥 Generate Dynamic Pose", variant="primary", size="lg")
            
            gr.Markdown("### 🔧 Advanced Refinement")
            refinement = gr.Textbox(
                label="Refinement command",
                placeholder="e.g., rotate, scale up, jump higher, spread arms, twist body",
                lines=1
            )
            
            refine_btn = gr.Button("✨ Apply Refinement", variant="secondary")
            
            gr.Examples(
                examples=[
                    "Ninja performing a flying kick",
                    "Basketball player doing a 360 dunk",
                    "Breakdancer doing a freeze",
                    "Gymnast doing a backflip",
                    "Martial artist in combat stance",
                    "Dancer leaping through the air",
                    "Rock climber reaching for a hold",
                    "Skateboarder doing a trick"
                ],
                inputs=prompt
            )
            
            gr.Markdown("""
            ### 💡 Pro Tips:
            - Use action verbs: jump, kick, flip, spin
            - Add intensity: extreme, dynamic, explosive
            - Specify sports/activities for better results
            - Combine with physics effects for realism
            """)
        
        with gr.Column(scale=1):
            pose_image = gr.Image(
                label="🎨 Generated Dynamic Pose",
                type="pil",
                height=512
            )
            
            with gr.Accordion("📊 OpenPose JSON Data", open=False):
                json_output = gr.Code(
                    language="json",
                    lines=15
                )
    
    # Event handlers
    generate_btn.click(
        fn=generate_pose,
        inputs=[prompt, use_llm, template, physics_effect],
        outputs=[pose_image, json_output, keypoints_state]
    )
    
    refine_btn.click(
        fn=refine_existing_pose,
        inputs=[refinement, keypoints_state],
        outputs=[pose_image, json_output, keypoints_state]
    )

# Launch
app.queue()
app.launch()