CoolFace
Apppublic

kchen707/wedding-bundle-builder

sourceHugging Facemitupdated 5mo agoView on Hugging Face
0likes
style_extraction.py89 linesDownload Raw Back to root
1"""2Vision-based style extraction from inspiration images.3Direct port of cell 31, with the OpenAI client switched to lazy.4"""5 6import base647from pathlib import Path8 9from clients import get_llm_client10from config import VISION_MODEL11from matching import _parse_json_response12 13 14VISION_STYLE_PROMPT = """You are a wedding style analyst. The user uploaded inspiration images for their wedding.15 16Look at all the images together and extract the visual style. Return ONLY a valid JSON object:17 18{19  "style_keywords": [<5-10 short style descriptors, e.g. "modern", "rustic", "bohemian", "minimalist", "romantic", "industrial", "vintage", "tropical", "elegant", "whimsical", "moody", "bright">],20  "color_palette": [<3-6 dominant colors, e.g. "blush pink", "sage green", "cream", "gold">],21  "setting_cues": [<2-4 venue/setting hints, e.g. "outdoor garden", "beach", "ballroom", "barn", "rooftop">],22  "floral_style": "<1 short phrase describing floral aesthetic, or null>",23  "overall_summary": "<1 sentence capturing the vibe>"24}25 26Rules:27- Focus on AESTHETIC cues, not specific objects28- If images are inconsistent, pick the dominant shared style29- Keep keywords SHORT (1-2 words each)30- Return ONLY the JSON object. No markdown fences, no preamble."""31 32 33def _encode_image_to_data_url(image_path):34    """Read an image file and return a base64 data URL."""35    path = Path(image_path)36    ext = path.suffix.lower().lstrip(".")37    mime_map = {"jpg": "jpeg", "jpeg": "jpeg", "png": "png",38                "gif": "gif", "webp": "webp"}39    mime = mime_map.get(ext, "jpeg")40    with open(path, "rb") as f:41        b64 = base64.b64encode(f.read()).decode("utf-8")42    return f"data:image/{mime};base64,{b64}"43 44 45def extract_style_from_images(image_paths, max_images=5):46    """47    Send up to max_images inspiration images to a vision model and return48    the extracted style dict. Returns None on failure.49    """50    if not image_paths:51        return None52 53    paths = [p for p in image_paths if p][:max_images]54    if not paths:55        return None56 57    content = [58        {"type": "text",59         "text": "Extract the wedding style from these inspiration images."},60    ]61    for p in paths:62        try:63            data_url = _encode_image_to_data_url(p)64            content.append({65                "type": "image_url",66                "image_url": {"url": data_url},67            })68        except Exception as e:69            print(f"Skipping {p}: {e}")70 71    if len(content) == 1:72        return None73 74    try:75        client = get_llm_client()76        resp = client.chat.completions.create(77            model=VISION_MODEL,78            messages=[79                {"role": "system", "content": VISION_STYLE_PROMPT},80                {"role": "user", "content": content},81            ],82            temperature=0.3,83        )84        raw = resp.choices[0].message.content85        return _parse_json_response(raw)86    except Exception as e:87        print(f"Vision extraction failed: {e}")88        return None89