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( "
" "Powered by OpenRouter & Nvidia Nemotron" "
" ) if __name__ == "__main__": demo.launch(share=False)