import gradio as gr
import easyocr
import cv2
import numpy as np
import re
import os

# Initialize EasyOCR Reader
reader = easyocr.Reader(['en'], gpu=False)

# Save processed image for inspection
DEBUG_SAVE_IMAGE = True  # Turn this off in production

# Preprocessing Function
def preprocess(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # Sharpen for deblurring
    sharpen_kernel = np.array([[-1, -1, -1],
                                [-1,  9, -1],
                                [-1, -1, -1]])
    sharp = cv2.filter2D(gray, -1, sharpen_kernel)
    
    # Adaptive Thresholding
    thresh = cv2.adaptiveThreshold(
        sharp, 255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV, 13, 2
    )

    # Dilate to enhance characters
    kernel = np.ones((2, 2), np.uint8)
    dilated = cv2.dilate(thresh, kernel, iterations=1)

    if DEBUG_SAVE_IMAGE:
        os.makedirs("debug", exist_ok=True)
        cv2.imwrite("debug/processed_debug.png", dilated)

    return dilated

# Extract weight with regex
def extract_weight(results):
    for result in results:
        text = result[0] if isinstance(result, (list, tuple)) else ""
        if isinstance(text, str):
            match = re.search(r'\b\d{2,4}(\.\d{1,2})?\s?(kg|lb|KG|LB)?\b', text)
            if match:
                return match.group()
    return None

# Main logic
def get_weight(img):
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

    # Upscale low-res images (distant capture)
    if img.shape[0] < 500:
        img = cv2.resize(img, None, fx=2.0, fy=2.0, interpolation=cv2.INTER_CUBIC)

    pre = preprocess(img)

    # Try OCR
    results = reader.readtext(pre)
    print("📝 OCR Results:", results)  # Log what OCR sees

    # Fallback: try raw image if nothing found
    weight = extract_weight(results)
    if not weight:
        print("🔁 Trying fallback: raw image")
        fallback_results = reader.readtext(img)
        print("📝 OCR Fallback:", fallback_results)
        weight = extract_weight(fallback_results)

    return f"Detected Weight: {weight if weight else 'Not found'}"

# Gradio UI
demo = gr.Interface(
    fn=get_weight,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=gr.Textbox(label="Detected Weight"),
    title="Auto Weight Logger (Smart OCR)",
    description="Extracts weight even from blurry or distant display images."
)

if __name__ == "__main__":
    demo.launch()