HairStyle.py

import os
import cv2
import base64
import requests
import json
import argparse
import re
import sys

# ---------------------------
# CONFIGURATION
# ---------------------------
API_URL = "https://openrouter.ai/api/v1/chat/completions"
MODEL_NAME = os.environ.get("OPENROUTER_MODEL", "nvidia/nemotron-nano-12b-v2-vl:free")

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 that instructs the model to return a JSON array
    containing exactly `n` hairstyle names (strings) chosen from the
    `MEN_HAIRSTYLES` list. The model must output ONLY the JSON array.
    """
    styles_list = ', '.join(MEN_HAIRSTYLES)
    return f"""
You are a professional MEN'S hairstylist.

TASK:
From the provided image, select the best {n} hairstyle names for this man.

RULES:
- Choose ONLY from this list: {styles_list}
- Output MUST be valid JSON: an array of {n} strings, for example ["Buzz Cut", "Fade"]
- Do NOT include explanations, reasons, or any extra text—ONLY the JSON array.
"""


def extract_styles_from_text(text, n):
    """Try to extract an ordered list of up to `n` hairstyles from text.
    First try to parse JSON; if that fails, search for known labels.
    """
    text = text.strip()
    # Try JSON
    try:
        parsed = json.loads(text)
        if isinstance(parsed, list):
            result = [s for s in parsed if isinstance(s, str) and s in MEN_HAIRSTYLES]
            return result[:n]
    except Exception:
        pass

    # Fallback: find known hairstyles in order of appearance
    found = []
    for m in re.finditer(r"\b(" + '|'.join(re.escape(s) for s in MEN_HAIRSTYLES) + r")\b", text, flags=re.IGNORECASE):
        name = m.group(0)
        # normalize to canonical casing from MEN_HAIRSTYLES
        for s in MEN_HAIRSTYLES:
            if name.lower() == s.lower() and s not in found:
                found.append(s)
                break
        if len(found) >= n:
            break
    return found[:n]


def main():
    parser = argparse.ArgumentParser(description="Get hairstyle recommendations via OpenRouter API")
    parser.add_argument("--image", "-i", help="Path to image file")
    parser.add_argument("--n", "-n", type=int, default=1, help="Number of styles to return (default 1)")
    args = parser.parse_args()

    image_path = args.image or input("📂 Enter full path to the image file: ").strip()
    if not image_path:
        print("No image path provided.")
        sys.exit(1)

    image = cv2.imread(image_path)
    if image is None:
        print("Failed to read image.")
        sys.exit(1)

    _, buffer = cv2.imencode(".png", image)
    image_b64 = base64.b64encode(buffer).decode("utf-8")

    api_key = "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"

    prompt = build_prompt(args.n)

    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.0,
        "max_tokens": 300
    }

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

    resp = requests.post(API_URL, headers=headers, json=payload, timeout=60)
    if resp.status_code != 200:
        print(f"API ERROR {resp.status_code}: {resp.text}")
        sys.exit(1)

    data = resp.json()
    try:
        message = data["choices"][0]["message"]
    except Exception:
        print("Unexpected API response format:", json.dumps(data))
        sys.exit(1)

    # message.content might be a string or a list; normalize to text
    text_content = ""
    content = message.get("content")
    if isinstance(content, str):
        text_content = content
    elif isinstance(content, list):
        parts = []
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                parts.append(item.get("text", ""))
            elif isinstance(item, str):
                parts.append(item)
        text_content = "\n".join(parts)
    else:
        text_content = str(content)

    styles = extract_styles_from_text(text_content, args.n)
    if not styles:
        # try reasoning_details fallback
        reasoning = message.get("reasoning_details") or []
        if reasoning and isinstance(reasoning, list):
            rd_text = reasoning[0].get("text", "")
            styles = extract_styles_from_text(rd_text, args.n)

    # Output just the recommended styles as a JSON array (no extra text)
    print(json.dumps(styles))


if __name__ == "__main__":
    main()