Update app.py

app.py

@@ -1,97 +1,113 @@
 import os
-import json
 import base64
+import json
+import google.generativeai as genai
 from flask import Flask, request, jsonify
-from groq import Groq
-from werkzeug.utils import secure_filename
 
 app = Flask(__name__)
-app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max file size
-app.config['UPLOAD_FOLDER'] = 'uploads'
 
-# Initialize Groq client with your API key
-client = Groq(api_key="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")
+# Configure Gemini
+genai.configure(api_key='AIzaSyDcYyq3w21iwipYn17wCAQo3AYWhUIGDSI')
+model = genai.GenerativeModel(model_name="gemini-1.5-pro")
 
-# Ensure upload folder exists
-os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
+# Configure upload settings
+UPLOAD_FOLDER = 'uploads'
+ALLOWED_EXTENSIONS = {'jpg', 'jpeg', 'png'}
 
-def encode_image_to_base64(image_path):
-    with open(image_path, "rb") as image_file:
-        return base64.b64encode(image_file.read()).decode("utf-8")
+if not os.path.exists(UPLOAD_FOLDER):
+    os.makedirs(UPLOAD_FOLDER)
 
-def extract_business_card_info(image_base64):
-    completion = client.chat.completions.create(
-        model="llama-3.2-11b-vision-preview",
-        messages=[
-            {
-                "role": "user",
-                "content": [
-                    {
-                        "type": "text",
-                        "text": """Extract info STRICTLY in this JSON format ONLY:
-{
-  "businessCard": {
-    "companyName": "",
-    "email": "",
-    "contactNumber": "",
-    "website": "",
-    "address": ""
-  }
-}
-NO TEXT BEFORE/AFTER THE JSON. NO MARKDOWN. ONLY VALID JSON."""
-                    },
-                    {
-                        "type": "image_url",
-                        "image_url": {
-                            "url": f"data:image/jpeg;base64,{image_base64}"
-                        }
-                    }
-                ]
-            }
-        ],
-        temperature=0.3,
-        max_tokens=1024,
-        top_p=1,
-        stream=False,
-    )
-    return completion.choices[0].message.content
+def allowed_file(filename):
+    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
 
-@app.route('/extract', methods=['POST'])
-def extract_info():
-    if 'image' not in request.files:
-        return jsonify({'error': 'No image file provided'}), 400
-    
-    file = request.files['image']
-    if file.filename == '':
-        return jsonify({'error': 'No selected file'}), 400
+def clean_response(response_text):
+    """Clean the response text to extract pure JSON."""
+    # Remove markdown code blocks
+    clean_text = response_text.replace('```json', '').replace('```', '').strip()
     
-    if file:
-        filename = secure_filename(file.filename)
-        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
-        file.save(filepath)
+    try:
+        # Parse and return clean JSON
+        json_data = json.loads(clean_text)
         
-        try:
-            # Convert image to base64
-            image_base64 = encode_image_to_base64(filepath)
+        # Clean up email and contact number fields (remove duplicates and format)
+        if 'businessCard' in json_data:
+            card = json_data['businessCard']
             
-            # Extract information
-            json_data = extract_business_card_info(image_base64)
+            # Clean email field
+            if card.get('email'):
+                emails = [e.strip() for e in card['email'].replace('\\n', ',').split(',')]
+                card['email'] = ', '.join(list(dict.fromkeys(emails)))  # Remove duplicates
             
-            # Clean and parse JSON
-            json_start = json_data.find('{')
-            json_end = json_data.rfind('}') + 1
-            clean_json = json_data[json_start:json_end]
-            parsed_json = json.loads(clean_json)
+            # Clean contact number field
+            if card.get('contactNumber'):
+                numbers = [n.strip() for n in card['contactNumber'].split(',')]
+                card['contactNumber'] = ', '.join(list(dict.fromkeys(numbers)))  # Remove duplicates
             
-            # Clean up the uploaded file
-            os.remove(filepath)
-            
-            return jsonify(parsed_json)
-            
-        except Exception as e:
-            if os.path.exists(filepath):
-                os.remove(filepath)
-            return jsonify({'error': str(e)}), 500
+            # Clean address field
+            if card.get('address'):
+                card['address'] = card['address'].replace('\\n', ' ').strip()
+                
+        return json_data
+    except json.JSONDecodeError as e:
+        raise Exception(f"Failed to parse JSON response: {str(e)}")
+
+def process_image(image_bytes):
+    base64_image = base64.b64encode(image_bytes).decode('utf-8')
+    prompt = """Extract info STRICTLY in this JSON format ONLY:
+    {
+      "businessCard": {
+        "companyName": "",
+        "email": "",
+        "contactNumber": "",
+        "website": "",
+        "address": ""
+      }
+    }
+    NO TEXT BEFORE/AFTER THE JSON. NO MARKDOWN. ONLY VALID JSON."""
+    
+    try:
+        response = model.generate_content(
+            [
+                {
+                    "mime_type": "image/jpeg",
+                    "data": base64_image
+                },
+                prompt
+            ]
+        )
+        # Clean and parse the response
+        return clean_response(response.text)
+    except Exception as e:
+        raise Exception(f"Error processing image: {str(e)}")
+
+@app.route('/process_card', methods=['POST'])
+def process_business_card():
+    # Check if file is present in request
+    if 'file' not in request.files:
+        return jsonify({'error': 'No file part'}), 400
+    
+    file = request.files['file']
+    if file.filename == '':
+        return jsonify({'error': 'No selected file'}), 400
+    
+    if not allowed_file(file.filename):
+        return jsonify({'error': 'File type not allowed. Use jpg, jpeg, or png'}), 400
+    
+    try:
+        # Read file bytes directly from the uploaded file
+        image_bytes = file.read()
+        
+        # Process the image and get clean JSON response
+        result = process_image(image_bytes)
+        
+        return jsonify(result), 200
+        
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+@app.route('/health', methods=['GET'])
+def health_check():
+    return jsonify({'status': 'healthy'}), 200
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860)
+    app.run(debug=True, port=7860)