receipt-api / main.py
rairo's picture
Update main.py
cde4684 verified
import os
import io
import json
import hashlib
from datetime import datetime
from PIL import Image
from flask import Flask, request, jsonify
from flask_cors import CORS
import google.generativeai as genai
app = Flask(__name__)
CORS(app) # Enable CORS for all routes
# Use the Gemini API key from the environment, or set it here for testing.
api_key = os.environ.get('Gemini', 'YOUR_GEMINI_API_KEY')
def configure_gemini(api_key):
genai.configure(api_key=api_key)
return genai.GenerativeModel('gemini-2.0-flash-thinking-exp')
def process_receipt(model, image):
prompt = (
"Analyze this image and determine if it's a receipt. If it is a receipt, extract:\n"
" - Total amount (as float)\n"
" - List of items purchased (array of strings)\n"
" - Date of transaction (DD/MM/YYYY format)\n"
" - Receipt number (as string)\n"
"Return JSON format with keys: is_receipt (boolean), total, items, date, receipt_number.\n"
"If not a receipt, return {\"is_receipt\": false}"
)
response = model.generate_content([prompt, image])
return response.text
@app.route('/process-receipt', methods=['POST'])
def process_receipt_endpoint():
try:
if 'receipt' not in request.files:
return jsonify({'error': 'No file uploaded'}), 400
file = request.files['receipt']
if file.filename == '':
return jsonify({'error': 'No file selected'}), 400
# Read file bytes and compute a hash (for duplicate checking or logging)
image_bytes = file.read()
file_hash = hashlib.md5(image_bytes).hexdigest()
# Open the image using Pillow
image = Image.open(io.BytesIO(image_bytes))
# Configure Gemini and process the receipt image
model = configure_gemini(api_key)
result_text = process_receipt(model, image)
# Attempt to extract JSON from the response text
json_start = result_text.find('{')
json_end = result_text.rfind('}')
if json_start == -1 or json_end == -1:
return jsonify({'error': 'Invalid response format', 'raw': result_text}), 500
json_str = result_text[json_start:json_end+1]
# Clean up any markdown formatting if necessary
json_str = json_str.replace('```json', '').replace('```', '')
data = json.loads(json_str)
# Optionally, add metadata to the response
data['file_hash'] = file_hash
data['timestamp'] = datetime.now().isoformat()
return jsonify(data)
except Exception as e:
return jsonify({'error': str(e)}), 500
if __name__ == '__main__':
# The server listens on all interfaces at port 7860.
app.run(debug=True, host="0.0.0.0", port=7860)