05/31 Phase 7: receipt ocr
This commit is contained in:
@@ -0,0 +1,227 @@
|
||||
"""
|
||||
OCR Service — extracts transaction data from receipt images using Groq vision.
|
||||
|
||||
Model: meta-llama/llama-4-scout-17b-16e-instruct (multimodal)
|
||||
Input: image file path or bytes
|
||||
Output: dict with amount, date, merchant, category_suggestion, notes
|
||||
"""
|
||||
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from datetime import date, datetime
|
||||
from flask import current_app
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
GROQ_VISION_MODEL = 'meta-llama/llama-4-scout-17b-16e-instruct'
|
||||
SUPPORTED_MIME_TYPES = {
|
||||
'.jpg': 'image/jpeg',
|
||||
'.jpeg': 'image/jpeg',
|
||||
'.png': 'image/png',
|
||||
'.gif': 'image/gif',
|
||||
'.webp': 'image/webp',
|
||||
}
|
||||
|
||||
RECEIPT_PROMPT = """You are a receipt data extractor. Analyze this receipt image and extract the following information.
|
||||
|
||||
Respond ONLY with a valid JSON object, no explanation, no markdown, no code fences. Just raw JSON.
|
||||
|
||||
Required format:
|
||||
{
|
||||
"amount": <total amount as a number, no currency symbols, e.g. 85.50>,
|
||||
"date": "<date in YYYY-MM-DD format, or null if not found>",
|
||||
"merchant": "<store/restaurant/vendor name, or null if not found>",
|
||||
"category_suggestion": "<one of: Food & Dining, Transport, Shopping, Health, Entertainment, Utilities, Housing, Education, Personal Care, Travel, Other>",
|
||||
"notes": "<brief description, e.g. 'Lunch at McDonald\\'s' or 'Grocery shopping'>"
|
||||
}
|
||||
|
||||
Rules:
|
||||
- amount: use the TOTAL/GRAND TOTAL amount, not subtotal. Numbers only, no commas.
|
||||
- date: format as YYYY-MM-DD. If only month/day visible, use current year.
|
||||
- merchant: the business name only, no address.
|
||||
- category_suggestion: pick the single most appropriate category from the list.
|
||||
- notes: keep under 60 characters.
|
||||
- If you cannot read the receipt clearly, still return your best guess with available data.
|
||||
- Never return null for amount if any number is visible."""
|
||||
|
||||
|
||||
def extract_from_file(file_path):
|
||||
"""
|
||||
Extract receipt data from an image file on disk.
|
||||
Returns dict or None on failure.
|
||||
"""
|
||||
ext = os.path.splitext(file_path)[1].lower()
|
||||
mime_type = SUPPORTED_MIME_TYPES.get(ext)
|
||||
|
||||
if not mime_type:
|
||||
log.warning(f'[ocr] Unsupported file type: {ext}')
|
||||
return _error_result(f'Unsupported file type: {ext}')
|
||||
|
||||
try:
|
||||
with open(file_path, 'rb') as f:
|
||||
image_bytes = f.read()
|
||||
return extract_from_bytes(image_bytes, mime_type)
|
||||
except FileNotFoundError:
|
||||
log.error(f'[ocr] File not found: {file_path}')
|
||||
return _error_result('Receipt file not found')
|
||||
except Exception as e:
|
||||
log.error(f'[ocr] File read error: {e}')
|
||||
return _error_result(str(e))
|
||||
|
||||
|
||||
def extract_from_bytes(image_bytes, mime_type='image/jpeg'):
|
||||
"""
|
||||
Extract receipt data from raw image bytes.
|
||||
Returns dict with keys: amount, date, merchant, category_suggestion, notes, error
|
||||
"""
|
||||
api_key = current_app.config.get('GROQ_API_KEY', '')
|
||||
if not api_key:
|
||||
return _error_result('GROQ_API_KEY not configured')
|
||||
|
||||
# Encode to base64
|
||||
b64 = base64.b64encode(image_bytes).decode('utf-8')
|
||||
data_url = f'data:{mime_type};base64,{b64}'
|
||||
|
||||
try:
|
||||
import requests
|
||||
resp = requests.post(
|
||||
'https://api.groq.com/openai/v1/chat/completions',
|
||||
headers={
|
||||
'Authorization': f'Bearer {api_key}',
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
json={
|
||||
'model': GROQ_VISION_MODEL,
|
||||
'messages': [
|
||||
{
|
||||
'role': 'user',
|
||||
'content': [
|
||||
{
|
||||
'type': 'image_url',
|
||||
'image_url': {'url': data_url},
|
||||
},
|
||||
{
|
||||
'type': 'text',
|
||||
'text': RECEIPT_PROMPT,
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
'max_tokens': 512,
|
||||
'temperature': 0.1, # low temp for consistent structured output
|
||||
},
|
||||
timeout=30,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
content = resp.json()['choices'][0]['message']['content'].strip()
|
||||
log.info(f'[ocr] Raw response: {content[:200]}')
|
||||
return _parse_response(content)
|
||||
|
||||
except requests.exceptions.HTTPError as e:
|
||||
if e.response.status_code == 400:
|
||||
# Model may not support this image — try to parse error
|
||||
try:
|
||||
err_body = e.response.json()
|
||||
msg = err_body.get('error', {}).get('message', str(e))
|
||||
except Exception:
|
||||
msg = str(e)
|
||||
log.error(f'[ocr] Groq 400 error: {msg}')
|
||||
return _error_result(f'Image processing failed: {msg[:100]}')
|
||||
elif e.response.status_code == 429:
|
||||
return _error_result('Rate limit reached. Try again shortly.')
|
||||
elif e.response.status_code == 401:
|
||||
return _error_result('Invalid Groq API key.')
|
||||
else:
|
||||
log.error(f'[ocr] HTTP {e.response.status_code}: {e}')
|
||||
return _error_result(f'API error ({e.response.status_code})')
|
||||
except Exception as e:
|
||||
log.error(f'[ocr] Unexpected error: {e}')
|
||||
return _error_result('OCR service temporarily unavailable')
|
||||
|
||||
|
||||
def _parse_response(content):
|
||||
"""Parse the JSON response from Groq into a clean dict."""
|
||||
# Strip any accidental markdown fences
|
||||
content = re.sub(r'^```(?:json)?\s*', '', content, flags=re.MULTILINE)
|
||||
content = re.sub(r'\s*```$', '', content, flags=re.MULTILINE)
|
||||
content = content.strip()
|
||||
|
||||
try:
|
||||
data = json.loads(content)
|
||||
except json.JSONDecodeError:
|
||||
# Try to extract JSON object from surrounding text
|
||||
match = re.search(r'\{.*\}', content, re.DOTALL)
|
||||
if match:
|
||||
try:
|
||||
data = json.loads(match.group())
|
||||
except json.JSONDecodeError:
|
||||
log.error(f'[ocr] Could not parse JSON: {content[:200]}')
|
||||
return _error_result('Could not parse OCR response')
|
||||
else:
|
||||
return _error_result('No JSON found in OCR response')
|
||||
|
||||
# Normalise amount
|
||||
amount = data.get('amount')
|
||||
if amount is not None:
|
||||
try:
|
||||
amount = float(str(amount).replace(',', '').strip())
|
||||
if amount <= 0:
|
||||
amount = None
|
||||
except (ValueError, TypeError):
|
||||
amount = None
|
||||
|
||||
# Normalise date
|
||||
raw_date = data.get('date')
|
||||
parsed_date = None
|
||||
if raw_date:
|
||||
for fmt in ('%Y-%m-%d', '%m/%d/%Y', '%d/%m/%Y', '%Y/%m/%d'):
|
||||
try:
|
||||
parsed_date = datetime.strptime(str(raw_date).strip(), fmt).date().isoformat()
|
||||
break
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
# Map category suggestion to our system category names
|
||||
cat_map = {
|
||||
'food & dining': 'Food & Dining',
|
||||
'food': 'Food & Dining',
|
||||
'dining': 'Food & Dining',
|
||||
'transport': 'Transport',
|
||||
'transportation':'Transport',
|
||||
'shopping': 'Shopping',
|
||||
'health': 'Health',
|
||||
'entertainment': 'Entertainment',
|
||||
'utilities': 'Utilities',
|
||||
'housing': 'Housing',
|
||||
'education': 'Education',
|
||||
'personal care': 'Personal Care',
|
||||
'travel': 'Travel',
|
||||
'other': 'Other',
|
||||
}
|
||||
raw_cat = str(data.get('category_suggestion', '')).lower().strip()
|
||||
category = cat_map.get(raw_cat, data.get('category_suggestion', 'Other'))
|
||||
|
||||
return {
|
||||
'amount': amount,
|
||||
'date': parsed_date or date.today().isoformat(),
|
||||
'merchant': data.get('merchant') or '',
|
||||
'category_suggestion': category,
|
||||
'notes': data.get('notes') or '',
|
||||
'error': None,
|
||||
'raw': data,
|
||||
}
|
||||
|
||||
|
||||
def _error_result(msg):
|
||||
return {
|
||||
'amount': None,
|
||||
'date': date.today().isoformat(),
|
||||
'merchant': '',
|
||||
'category_suggestion': 'Other',
|
||||
'notes': '',
|
||||
'error': msg,
|
||||
'raw': {},
|
||||
}
|
||||
Reference in New Issue
Block a user