import gradio as gr
from IndicPhotoOCR.ocr import OCR
import torch
import cv2
import numpy as np

def visualize_detection(image, detections):
    for box in detections:
        points = np.array(box, np.int32)
        x_min = np.min(points[:, 0])
        y_min = np.min(points[:, 1])
        x_max = np.max(points[:, 0])
        y_max = np.max(points[:, 1])
        cv2.rectangle(image, (x_min, y_min), (x_max, y_max), color=(0, 255, 0), thickness=3)
    return image

device = "cuda" if torch.cuda.is_available() else "cpu"

ocr = OCR(device)
def process_image(image_path, identifier_lang):
    global ocr
    if identifier_lang != ocr.indentifier_lang:
        del ocr
        ocr = OCR(device,identifier_lang)
    ocr_data,ocr_detect = ocr.ocr(image_path)
    return '\n'.join([' '.join(line) for line in ocr_data]),visualize_detection(image_path, ocr_detect)

gr.Interface(process_image,[gr.Image(type="pil"),gr.Dropdown(["hindi", "assamese", "bengali", "gujarati", "kannada", "malayalam","odia", "punjabi", "tamil", "telugu", "auto"], label="Language")],[gr.Textbox(label="Recognized Text"),gr.Image(type="pil")]).launch()