`python-matcher/scripts/embed_handbags.py`

129 lines
"""
Batch script to embed handbag catalog images
Processes all handbags with embedding IS NULL
"""
import os
import sys
import io
import psycopg2
import requests
from PIL import Image
from datetime import datetime

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
from models.handbag_embedder import image_to_vec, embed_batch

DB_DSN = os.environ.get("DATABASE_URL", "postgresql://localhost/handbags")

def get_conn():
    """Get database connection"""
    return psycopg2.connect(DB_DSN)


def fetch_unembedded_handbags(conn, limit=500):
    """Fetch handbags that need embeddings"""
    with conn.cursor() as cur:
        cur.execute("""
            SELECT id, sku, image_url, brand, model_name
            FROM handbags
            WHERE embedding IS NULL
              AND is_active = true
              AND image_url IS NOT NULL
            ORDER BY id
            LIMIT %s;
        """, (limit,))
        return cur.fetchall()


def update_embedding(conn, row_id, vec):
    """Update handbag with embedding vector"""
    with conn.cursor() as cur:
        # pgvector expects a string like '[0.1, 0.2, ...]'
        vec_str = "[" + ",".join(str(x) for x in vec) + "]"
        cur.execute("""
            UPDATE handbags
            SET embedding = %s,
                updated_at = CURRENT_TIMESTAMP
            WHERE id = %s;
        """, (vec_str, row_id))


def download_image(url, timeout=10):
    """Download and open image from URL"""
    headers = {
        'User-Agent': 'Mozilla/5.0 (compatible; HandbagEmbedder/1.0)'
    }
    resp = requests.get(url, timeout=timeout, headers=headers)
    resp.raise_for_status()
    img = Image.open(io.BytesIO(resp.content)).convert("RGB")
    return img


def main():
    """Main embedding loop"""
    print("🚀 Starting handbag embedding process...")
    print(f"📊 Database: {DB_DSN}\n")

    conn = get_conn()
    conn.autocommit = False

    total_processed = 0
    total_errors = 0

    try:
        while True:
            # Fetch batch of unembedded handbags
            rows = fetch_unembedded_handbags(conn, limit=100)

            if not rows:
                print("\n✅ No more handbags to embed!")
                break

            print(f"\n📦 Processing batch of {len(rows)} handbags...")

            for row_id, sku, image_url, brand, model in rows:
                try:
                    # Download image
                    img = download_image(image_url)

                    # Generate embedding
                    vec = image_to_vec(img)

                    # Update database
                    update_embedding(conn, row_id, vec)

                    total_processed += 1
                    print(f"✓ [{total_processed}] {brand} {model} (SKU: {sku})")

                except Exception as e:
                    total_errors += 1
                    print(f"✗ Error on {sku}: {e}")
                    # Continue with next item

            # Commit batch
            conn.commit()
            print(f"💾 Committed batch ({total_processed} total, {total_errors} errors)")

        # Final stats
        print(f"\n" + "="*60)
        print(f"🎉 Embedding complete!")
        print(f"   ✅ Processed: {total_processed}")
        print(f"   ❌ Errors: {total_errors}")
        print(f"   📊 Success rate: {(total_processed/(total_processed+total_errors)*100):.1f}%")
        print("="*60)

    except KeyboardInterrupt:
        print("\n⚠️  Interrupted by user")
        conn.rollback()
    except Exception as e:
        print(f"\n❌ Fatal error: {e}")
        conn.rollback()
        raise
    finally:
        conn.close()


if __name__ == "__main__":
    main()