Spaces:
Running
Running
| import gradio as gr | |
| from PIL import Image | |
| from datetime import datetime | |
| import pytz | |
| from ocr_engine import extract_weight, extract_unit_from_text | |
| from simple_salesforce import Salesforce | |
| import base64 | |
| import re | |
| import os | |
| # Salesforce credentials | |
| SF_USERNAME = "[email protected]" | |
| SF_PASSWORD = "autoweight@32" | |
| SF_TOKEN = "UgiHKWT0aoZRX9gvTYDjAiRY" | |
| # Connect to Salesforce | |
| sf = Salesforce(username=SF_USERNAME, password=SF_PASSWORD, security_token=SF_TOKEN) | |
| def restore_decimal(text): | |
| if re.fullmatch(r"\d{5}", text): | |
| return f"{text[:2]}.{text[2:]}" | |
| elif re.fullmatch(r"\d{4}", text): | |
| return f"{text[:2]}.{text[2:]}" | |
| return text | |
| def process_image(image): | |
| if image is None: | |
| return "❌ No image provided", "", None, gr.update(visible=True) | |
| try: | |
| result = extract_weight(image) | |
| weight, raw_text = result if isinstance(result, tuple) else (result, "") | |
| print("🧠 Final OCR Result:", weight) | |
| print("🔤 OCR Raw Text:", raw_text) | |
| ist = pytz.timezone('Asia/Kolkata') | |
| timestamp = datetime.now(ist).strftime("%Y-%m-%d %H:%M:%S IST") | |
| if not weight or (isinstance(weight, str) and weight.startswith("Error")): | |
| return f"❌ OCR Error: {weight}", "", image, gr.update(visible=True) | |
| match = re.search(r'(\d{1,3}\.\d{1,3})\s*(kg|g)?', weight) | |
| if match: | |
| numeric_value = float(match.group(1)) | |
| unit = match.group(2) if match.group(2) else extract_unit_from_text(raw_text) | |
| else: | |
| cleaned = re.sub(r"[^\d]", "", weight) | |
| decimal_fixed = restore_decimal(cleaned) | |
| try: | |
| numeric_value = float(decimal_fixed) | |
| unit = extract_unit_from_text(raw_text) | |
| except: | |
| return f"❌ Could not extract number | OCR: {weight}", "", image, gr.update(visible=True) | |
| image_path = "snapshot.jpg" | |
| image.save(image_path) | |
| record = sf.Weight_Log__c.create({ | |
| "Captured_Weight__c": numeric_value, | |
| "Captured_Unit__c": unit, | |
| "Captured_At__c": datetime.now(ist).isoformat(), | |
| "Device_ID__c": "DEVICE-001", | |
| "Status__c": "Confirmed" | |
| }) | |
| with open(image_path, "rb") as f: | |
| encoded_image = base64.b64encode(f.read()).decode("utf-8") | |
| content = sf.ContentVersion.create({ | |
| "Title": f"Snapshot_{timestamp}", | |
| "PathOnClient": "snapshot.jpg", | |
| "VersionData": encoded_image | |
| }) | |
| content_id = sf.query( | |
| f"SELECT ContentDocumentId FROM ContentVersion WHERE Id = '{content['id']}'" | |
| )['records'][0]['ContentDocumentId'] | |
| sf.ContentDocumentLink.create({ | |
| "ContentDocumentId": content_id, | |
| "LinkedEntityId": record['id'], | |
| "ShareType": "V", | |
| "Visibility": "AllUsers" | |
| }) | |
| return f"{numeric_value} {unit}", timestamp, image, gr.update(visible=False) | |
| except Exception as e: | |
| return f"Error: {str(e)}", "", None, gr.update(visible=True) | |
| with gr.Blocks(css=".gr-button {background-color: #2e7d32 !important; color: white !important;}" ) as demo: | |
| gr.Markdown(""" | |
| <h1 style='text-align: center; color: #2e7d32;'>📷 Auto Weight Logger</h1> | |
| <p style='text-align: center;'>Upload or capture a digital weight image. Detects weight using AI OCR and logs it into Salesforce.</p> | |
| <hr style='border: 1px solid #ddd;' /> | |
| """) | |
| with gr.Row(): | |
| image_input = gr.Image(type="pil", label="📁 Upload or Capture Image") | |
| detect_btn = gr.Button("🚀 Detect Weight") | |
| with gr.Row(): | |
| weight_out = gr.Textbox(label="📦 Detected Weight", placeholder="e.g., 75.5 kg", show_copy_button=True) | |
| time_out = gr.Textbox(label="🕒 Captured At (IST)", placeholder="e.g., 2025-07-01 12:00:00") | |
| snapshot = gr.Image(label="📸 Snapshot Preview") | |
| retake_btn = gr.Button("🔁 Retake / Try Again", visible=False) | |
| detect_btn.click(fn=process_image, inputs=image_input, outputs=[weight_out, time_out, snapshot, retake_btn]) | |
| retake_btn.click(fn=lambda: ("", "", None, gr.update(visible=False)), | |
| inputs=[], outputs=[weight_out, time_out, snapshot, retake_btn]) | |
| if __name__ == "__main__": | |
| demo.launch() | |