import gradio as gr
from mmocr.apis import MMOCR
import cv2
import os

# Initialize MMOCR once
ocr = MMOCR(det='DB_r18', recog='SAR', device='cpu')  # use 'cuda:0' if GPU is available

def detect_weight(image):
    results = ocr.readtext(image, output=None)
    
    detected_texts = [item['text'] for item in results[0]['result']]
    
    # Search for weight-like patterns like "25.5 kg", "100kg", etc.
    weight = "Not detected"
    for text in detected_texts:
        if 'kg' in text.lower():
            weight = text
            break

    return f"Detected Weight: {weight}"

# Gradio UI
interface = gr.Interface(
    fn=detect_weight,
    inputs=gr.Image(type="filepath", label="Upload Weight Image"),
    outputs=gr.Textbox(label="OCR Output"),
    title="Weight Detector using MMOCR",
    description="Upload an image with weight text (e.g., '25 kg'), and this app will detect it."
)

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