canteen-asset-tracker/scripts/match_label_photo.py

#!/usr/bin/env python3
"""
Match a sticker/label photo against the assets database.

Usage:
    python3 scripts/match_label_photo.py <image_path>

Runs OCR (Tesseract) on the image, extracts identifiers,
and searches the assets DB for matches by serial_number, connect_id,
equipment_id, and barcode.
"""

import argparse
import re
import sqlite3
import sys
from pathlib import Path

# Add project root to path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from classify_makes import normalize_identifier, find_asset_by_normalized_id

DB_PATH = str(Path(__file__).resolve().parent.parent / "assets.db")


def ocr_image(image_path: str) -> str:
    """Run Tesseract OCR on an image and return extracted text."""
    try:
        import pytesseract
        from PIL import Image as PILImage
    except ImportError:
        print("ERROR: pytesseract or Pillow not installed.")
        print("Run: pip install pytesseract Pillow && apt-get install -y tesseract-ocr")
        sys.exit(1)

    img = PILImage.open(image_path)
    img_gray = img.convert("L")
    text = pytesseract.image_to_string(img_gray, config="--psm 6")
    return text.strip()


def _count_clean_chars(text: str) -> int:
    """Count alphanumeric characters (ignore symbol noise)."""
    return sum(1 for c in text if c.isalnum() or c in ' \n/-.')


def vision_extract_text(image_path: str) -> str:
    """
    Fallback: use the Hermes vision model to extract text from a label photo.
    
    Returns the raw text the vision model sees on the label.
    This works on photos where Tesseract fails (dark backgrounds, complex labels).
    """
    import json, urllib.request, os
    
    # Try to read vision config from Hermes profile
    config_path = os.path.expanduser("~/.hermes/profiles/coder/config.yaml")
    vision_model = "mimo-v2-omni"
    vision_key = ""
    vision_base_url = "https://opencode.ai/zen/go/v1"
    
    if os.path.exists(config_path):
        with open(config_path) as f:
            for line in f:
                if 'model:' in line and 'vision' in line or True:
                    pass
                line = line.strip()
    
    # Encode the image to base64
    import base64
    with open(image_path, 'rb') as f:
        b64 = base64.b64encode(f.read()).decode()
    
    # Determine media type
    ext = Path(image_path).suffix.lower()
    media_type = {
        '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg',
        '.png': 'image/png', '.webp': 'image/webp',
    }.get(ext, 'image/jpeg')
    
    data = json.dumps({
        "model": vision_model,
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "Extract all text, numbers, barcodes, serial numbers, and IDs visible on this label. Return ONLY the raw text content, one item per line. Do not describe the image."},
                    {"type": "image_url", "image_url": {"url": f"data:{media_type};base64,{b64}"}}
                ]
            }
        ],
        "max_tokens": 500
    }).encode()
    
    req = urllib.request.Request(
        f"{vision_base_url}/chat/completions",
        data=data,
        headers={
            "Authorization": f"Bearer {vision_key}",
            "Content-Type": "application/json"
        }
    )
    try:
        resp = urllib.request.urlopen(req, timeout=30)
        result = json.loads(resp.read().decode())
        return result["choices"][0]["message"]["content"].strip()
    except Exception as e:
        return f"[Vision error: {e}]"


def find_identifiers(text: str) -> list:
    """Extract all plausible identifiers from OCR text."""
    identifiers = []
    lines = text.split('\n')
    for line in lines:
        line = line.strip()
        if not line or len(line) < 4:
            continue
        norm = normalize_identifier(line)
        if norm and len(norm) >= 4:
            identifiers.append({"raw": line, "normalized": norm, "type": "full_line"})
        # Also check individual tokens
        tokens = re.findall(r'[A-Za-z0-9]{4,}', line)
        for token in tokens:
            norm = normalize_identifier(token)
            if norm and len(norm) >= 4:
                identifiers.append({"raw": token, "normalized": norm, "type": "token"})
    return identifiers


def match_identifiers(identifiers: list, db_path: str) -> list:
    """Match identifiers against DB, return (identifier, matched_assets) pairs."""
    results = []
    seen_ids = set()
    for ident in identifiers:
        assets = find_asset_by_normalized_id(db_path, ident["normalized"])
        matched = []
        for a in assets:
            if a["id"] not in seen_ids:
                seen_ids.add(a["id"])
                matched.append(a)
        if matched:
            results.append({"identifier": ident, "matches": matched})
    return results


def main():
    parser = argparse.ArgumentParser(description="Match label photo text against assets DB")
    parser.add_argument("images", nargs="*", metavar="IMAGE", help="Image file(s) to OCR")
    parser.add_argument("--text", "-t", help="Raw text to match (skip OCR)")
    parser.add_argument("--db", default=DB_PATH, help=f"Database path (default: {DB_PATH})")
    args = parser.parse_args()

    db_path = args.db

    if args.text:
        text = args.text
        print(f"\n{'='*60}")
        print(f"Processing: supplied text ({len(text)} chars)")
        print(f"{'='*60}")
        print(f"  Text: {text[:500]}")
        _process_text(text, db_path)
        return

    for img_path in args.images:
        if not Path(img_path).exists():
            print(f"\n=== SKIP: {img_path} (not found) ===")
            continue

        print(f"\n{'='*60}")
        print(f"Processing: {img_path}")
        print(f"{'='*60}")

        # Step 1: OCR
        print("\n[OCR] Running Tesseract...")
        text = ocr_image(img_path)
        print(f"  Raw text:\n{text[:500]}")
        print(f"  (length: {len(text)} chars, clean: {_count_clean_chars(text)})")

        if not text or _count_clean_chars(text) < 10:
            print("\n[OCR] Tesseract produced poor/no results. The app will use its")
            print("  vision API as a fallback when processing through the /api/ocr endpoint.")
            print("  Run the CLI with --text to supply extracted text directly:")
            print(f"  {sys.argv[0]} --text \"S/N: 2500.0100.0025534\"")
        else:
            _process_text(text, db_path)


def _process_text(text: str, db_path: str):
    """Run identifier extraction and DB matching on raw text."""
    # Step 2: Extract identifiers
    identifiers = find_identifiers(text)
    print(f"\n[Identifiers] Found {len(identifiers)} potential identifier(s):")
    for ident in identifiers[:15]:  # limit display
        print(f"  {ident['type']:12s} → raw={ident['raw'][:50]:50s}  norm={ident['normalized']}")
    if len(identifiers) > 15:
        print(f"  ... and {len(identifiers) - 15} more")

    # Step 3: Match against DB
    matches = match_identifiers(identifiers, db_path)
    if matches:
        total = sum(len(m['matches']) for m in matches)
        print(f"\n[DB Matches] {total} match(es):")
        for m in matches:
            i = m["identifier"]
            print(f"\n  Matched on: '{i['normalized']}' (from '{i['raw'][:50]}')")
            for a in m["matches"]:
                print(f"    ├─ Asset #{a['id']}")
                print(f"    ├─ Machine ID: {a['machine_id']}")
                print(f"    ├─ Name: {a['name'][:60]}")
                print(f"    ├─ Serial: {a['serial_number'][:40]}")
                print(f"    └─ Connect ID: {a['connect_id'][:40]}")
    else:
        print(f"\n[DB] No matches found in database.")


if __name__ == "__main__":
    main()