app.py

import gradio as gr
import cv2
import base64
import requests
import json
import numpy as np
from PIL import Image
import re
import io

# ---------------------------
# CONFIGURATION
# ---------------------------
API_URL = "https://openrouter.ai/api/v1/chat/completions"
MODEL_NAME = "nvidia/nemotron-nano-12b-v2-vl:free"
API_KEY = "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"

MEN_HAIRSTYLES = [
    "Buzz Cut", "Crew Cut", "Fade", "Undercut",
    "Slick Back", "Side Part", "Quiff", "Pompadour",
    "French Crop", "Textured Fringe"
]


def build_prompt(n):
    """Builds a prompt for hairstyle recommendations with gender validation"""
    styles_list = ', '.join(MEN_HAIRSTYLES)

    return f"""
You are a professional hairstylist analyzing a SINGLE PERSON in an image.

STEP 1 — GENDER CHECK:
- First, determine whether the person in the image is MALE or FEMALE.
- If the person is FEMALE, STOP immediately and return EXACTLY this JSON:
["Please upload a male photo for hairstyle recommendations"]

STEP 2 — MALE ONLY:
- If the person is MALE, continue with the task below.

TASK (MALE ONLY):
Select exactly {n} hairstyle names that would suit this man.

ALLOWED HAIRSTYLES (choose ONLY from this list):
{styles_list}

OUTPUT FORMAT RULES (STRICT):
- Return ONLY a JSON array
- If MALE → exactly {n} hairstyle names
- If FEMALE → the exact message shown above
- No explanations
- No extra text
- No numbering
- No markdown
- No line breaks before or after the JSON

EXAMPLES:
MALE:
["Buzz Cut", "Fade", "Crew Cut"]

FEMALE:
["Please upload a male photo for hairstyle recommendations"]
"""


def extract_styles_from_text(text, n):
    """Extract hairstyle names from text response"""
    text = text.strip()
    
    print(f"[DEBUG] Extracting styles from text length: {len(text)}")
    print(f"[DEBUG] Text preview: {text[:300]}")
    
    # Try JSON parsing first
    try:
        parsed = json.loads(text)
        if isinstance(parsed, list):
            result = [s for s in parsed if isinstance(s, str) and s in MEN_HAIRSTYLES]
            if result:
                print(f"[DEBUG] Found {len(result)} styles from JSON")
                return result[:n]
    except Exception as e:
        print(f"[DEBUG] JSON parse failed: {e}")
        pass

    # Strategy 1: Look for ["style1", "style2", "style3"] pattern
    json_pattern_match = re.search(r'\[\"([^\"]+)\"(?:,\s*\"([^\"]+)\")*\]', text)
    if json_pattern_match:
        # Extract all quoted items
        items = re.findall(r'\"([^\"]+)\"', json_pattern_match.group(0))
        found = [item for item in items if item in MEN_HAIRSTYLES]
        if found:
            print(f"[DEBUG] Found {len(found)} styles from JSON pattern")
            return found[:n]

    # Strategy 2: Try extracting from common sentence patterns
    lines = text.split('\n')
    found = []
    
    for line in lines:
        line = line.strip()
        # Remove common prefixes
        line = re.sub(r'^[-•*]\s*', '', line)
        line = re.sub(r'^\d+\.\s*', '', line)
        
        # Try exact match first
        if line in MEN_HAIRSTYLES:
            if line not in found:
                found.append(line)
                if len(found) >= n:
                    break
            continue
        
        # Try case-insensitive match
        for style in MEN_HAIRSTYLES:
            if line.lower() == style.lower() and style not in found:
                found.append(style)
                if len(found) >= n:
                    break
        
        if len(found) >= n:
            break
    
    # Strategy 3: Look for ALL mentions of hairstyles in text
    if len(found) < n:
        print(f"[DEBUG] Found {len(found)}, searching for more...")
        for style in MEN_HAIRSTYLES:
            # Look for the style name as a whole word
            pattern = r'\b' + re.escape(style) + r'\b'
            if re.search(pattern, text, re.IGNORECASE) and style not in found:
                found.append(style)
                print(f"[DEBUG] Found style: {style}")
                if len(found) >= n:
                    break
    
    print(f"[DEBUG] Extraction found {len(found)} styles: {found}")
    return found[:n]


def get_hairstyle_recommendations(image, num_styles):
    """Call OpenRouter API to get hairstyle recommendations"""
    if image is None:
        return "❌ Please upload an image first"
    
    try:
        # Convert to pillow image if needed
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image.astype('uint8'))
        elif not isinstance(image, Image.Image):
            image = Image.fromarray(np.array(image).astype('uint8'))
        
        # Ensure RGB mode
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        # Save to bytes and encode as base64
        buffer = io.BytesIO()
        image.save(buffer, format='PNG')
        image_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")

        prompt = build_prompt(num_styles)

        payload = {
            "model": MODEL_NAME,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {"type": "image_url", "image_url": f"data:image/png;base64,{image_b64}"}
                    ]
                }
            ],
            "temperature": 0.7,
            "max_tokens": 1000,
            "top_p": 0.9
        }

        headers = {
            "Authorization": f"Bearer {API_KEY}",
            "Content-Type": "application/json",
            "Accept": "application/json"
        }

        # Debug: Show request info
        print(f"[DEBUG] Sending request to {API_URL}")
        print(f"[DEBUG] Model: {MODEL_NAME}")
        print(f"[DEBUG] Image B64 length: {len(image_b64)}")

        resp = requests.post(API_URL, headers=headers, json=payload, timeout=60)
        
        print(f"[DEBUG] Response status: {resp.status_code}")
        print(f"[DEBUG] Response body: {resp.text[:500]}")
        
        if resp.status_code != 200:
            return f"❌ API ERROR {resp.status_code}\n\nResponse: {resp.text[:300]}"

        data = resp.json()
        
        # Check if response has choices
        if "choices" not in data:
            return f"❌ Invalid API response: {json.dumps(data)[:300]}"
        
        message = data["choices"][0]["message"]
        
        # Extract text content - check both regular content and reasoning
        content = message.get("content", "").strip()
        reasoning = message.get("reasoning", "").strip()
        
        # If content is empty, try reasoning
        text_content = content if content else reasoning
        
        # If still empty, try reasoning_details
        if not text_content:
            reasoning_details = message.get("reasoning_details", [])
            if reasoning_details and isinstance(reasoning_details, list) and len(reasoning_details) > 0:
                text_content = reasoning_details[0].get("text", "")
        
        print(f"[DEBUG] Content: {content[:100] if content else 'EMPTY'}")
        print(f"[DEBUG] Reasoning: {reasoning[:100] if reasoning else 'EMPTY'}")
        print(f"[DEBUG] Using text_content from: {'content' if content else 'reasoning'}")
        print(f"[DEBUG] Extracted text preview: {text_content[:300]}")

        styles = extract_styles_from_text(text_content, num_styles)
        
        if not styles:
            return f"⚠️ Could not extract recommendations.\n\nRaw API Response:\n{text_content}"
        
        result = "✅ Recommended Hairstyles:\n\n"
        for i, style in enumerate(styles, 1):
            result += f"{i}. {style}\n"
        
        return result
    
    except Exception as e:
        import traceback
        error_trace = traceback.format_exc()
        return f"❌ Error: {str(e)}\n\n{error_trace}"


# ---------------------------
# GRADIO INTERFACE
# ---------------------------
with gr.Blocks(title="Men's Hairstyle Recommender") as demo:
    gr.Markdown("# 💇‍♂️ Men's Hairstyle Recommender")
    gr.Markdown("Upload a photo of a man to get personalized hairstyle recommendations!")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Upload Image")
            image_input = gr.Image(
                label="Choose Image",
                type="pil",
                scale=1
            )
            
        with gr.Column():
            gr.Markdown("### Settings & Results")
            num_styles = gr.Slider(
                minimum=1,
                maximum=len(MEN_HAIRSTYLES),
                value=3,
                step=1,
                label="Number of Recommendations"
            )
            
            output = gr.Textbox(
                label="Recommendations",
                lines=10,
                interactive=False
            )
    
    # Button to get recommendations
    submit_btn = gr.Button("🎯 Get Recommendations", scale=1)
    submit_btn.click(
        fn=get_hairstyle_recommendations,
        inputs=[image_input, num_styles],
        outputs=output
    )
    
    # Sidebar info
    with gr.Accordion("Available Hairstyles"):
        hairstyles_text = "\n".join([f"• {style}" for style in MEN_HAIRSTYLES])
        gr.Markdown(hairstyles_text)
    
    gr.Markdown("---")
    gr.Markdown(
        "<div style='text-align: center;'>"
        "Powered by OpenRouter & Nvidia Nemotron"
        "</div>"
    )


if __name__ == "__main__":
    demo.launch(share=False)