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

ocr = PaddleOCR(use_angle_cls=True, lang='en')  # Better for clean numbers

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

    # Convert to OpenCV format
    image_np = np.array(image)

    # Preprocessing: Convert to grayscale, enhance contrast
    gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
    gray = cv2.equalizeHist(gray)
    gray = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)  # Convert back to 3 channel for PaddleOCR

    # Run OCR
    result = ocr.ocr(gray, cls=True)

    best_match = None
    best_conf = 0

    for line in result:
        for box in line:
            text = box[1][0]
            conf = box[1][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

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 weight scale image. The app will detect and log the weight automatically."
).launch()