import gradio as gr
from PIL import Image
import torch
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from datetime import datetime
import pytz

# Load the model and processor from Hugging Face
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-stage1")
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-stage1")

# Function to detect weight text from image
def detect_weight(image):
    try:
        # Preprocess image
        pixel_values = processor(images=image, return_tensors="pt").pixel_values
        # Run model
        generated_ids = model.generate(pixel_values)
        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

        # Try to extract weight-like number
        import re
        match = re.search(r"(\d{1,3}(\.\d{1,2})?)", generated_text)
        weight = match.group(1) if match else "Not detected"

        # Get IST time
        ist = pytz.timezone('Asia/Kolkata')
        current_time = datetime.now(ist).strftime("%Y-%m-%d %H:%M:%S")

        return f"Weight: {weight} kg\nCaptured At: {current_time}", image

    except Exception as e:
        return f"Error: {str(e)}", image

# Gradio UI
interface = gr.Interface(
    fn=detect_weight,
    inputs=gr.Image(type="pil", label="Upload or Capture Image"),
    outputs=[gr.Textbox(label="Weight Info"), gr.Image(label="Snapshot")],
    title="⚖️ Auto Weight Detector (No Tesseract)",
    description="Detects weight from digital scale image using AI-based OCR (no Tesseract)."
)

interface.launch()