import gradio as gr
import cv2
import numpy as np
from paddleocr import PaddleOCR
from datetime import datetime
import re
from PIL import Image

# Initialize PaddleOCR once
ocr = PaddleOCR(use_angle_cls=True, lang='en')  # English OCR model

def detect_weight(image):
    try:
        if image is None:
            return "No image uploaded", "N/A", None

        # Convert PIL image to RGB numpy array
        image = image.convert("RGB")
        image_np = np.array(image)

        # Convert to grayscale and enhance contrast
        gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
        enhanced = cv2.equalizeHist(gray)
        rgb_image = cv2.cvtColor(enhanced, cv2.COLOR_GRAY2RGB)

        # Perform OCR
        result = ocr.ocr(rgb_image, cls=True)

        best_match = None
        best_conf = 0

        for line in result:
            for box in line:
                text, conf = box[1]
                match = re.search(r"\d+\.\d+", text)
                if match and conf > best_conf:
                    best_match = match.group()
                    best_conf = conf

        if best_match:
            now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            return f"Weight: {best_match} kg (Confidence: {round(best_conf * 100, 2)}%)", now, image
        else:
            return "No weight detected kg (Confidence: 0.0%)", "N/A", image

    except Exception as e:
        # Show error in output instead of crashing
        return f"Error: {str(e)}", "Error", None

# Gradio UI
gr.Interface(
    fn=detect_weight,
    inputs=gr.Image(type="pil", label="Upload or Capture Weight Image"),
    outputs=[
        gr.Text(label="Detected Weight"),
        gr.Text(label="Captured At (IST)"),
        gr.Image(label="Snapshot")
    ],
    title="Auto Weight Logger",
    description="Upload or capture a digital scale image. This app detects the weight automatically using AI."
).launch()