import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO
from PIL import Image

# Load YOLO models
try:
    front_model = YOLO("front_cnic_model.pt")
    back_model = YOLO("back_cnic_model.pt")
except Exception as e:
    print(f"Error loading models: {e}")

def preprocess_image(image):
    """Convert PIL Image to OpenCV format."""
    try:
        img = np.array(image)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        return img
    except Exception as e:
        return None, f"Error preprocessing image: {e}"

def extract_regions(image, boxes, model_names):
    """Extract regions of interest (ROIs) by label from YOLO detections."""
    results = {}
    for box in boxes:
        try:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            cls_name = model_names[int(box.cls)]
            roi = image[y1:y2, x1:x2]
            results[cls_name] = f"Detected region at [{x1}, {y1}, {x2}, {y2}]"
        except Exception as e:
            results[cls_name] = f"Region extraction error: {e}"
    return results

def process_front_cnic(image):
    """Process front CNIC image to extract fields."""
    if image is None:
        return {"error": "No image provided"}
    
    img = preprocess_image(image)
    if isinstance(img, tuple):
        return {"error": img[1]}
    
    try:
        results = front_model.predict(img, conf=0.5)
        extracted_info = extract_regions(img, results[0].boxes, front_model.names)
        return extracted_info if extracted_info else {"error": "No objects detected"}
    except Exception as e:
        return {"error": f"Front CNIC processing error: {e}"}

def process_back_cnic(image):
    """Process back CNIC image to extract ROIs for Barcode and CNIC number."""
    if image is None:
        return {"error": "No image provided"}
    
    img = preprocess_image(image)
    if isinstance(img, tuple):
        return {"error": img[1]}
    
    try:
        results = back_model.predict(img, conf=0.5)
        extracted_info = extract_regions(img, results[0].boxes, back_model.names)
        return extracted_info if extracted_info else {"error": "No objects detected"}
    except Exception as e:
        return {"error": f"Back CNIC processing error: {e}"}

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # CNIC Detection (YOLO-only)
        Upload front and back CNIC images to detect regions of interest.
        - **Front CNIC**: Detects regions such as Name, DOB, Father's Name, etc.
        - **Back CNIC**: Detects Barcode and CNIC number regions (no OCR).
        """
    )
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("## Front CNIC")
            front_input = gr.Image(type="pil", label="Upload Front CNIC Image")
            front_button = gr.Button("Process Front CNIC")
            front_output = gr.JSON(label="Detected Regions")
        
        with gr.Column():
            gr.Markdown("## Back CNIC")
            back_input = gr.Image(type="pil", label="Upload Back CNIC Image")
            back_button = gr.Button("Process Back CNIC")
            back_output = gr.JSON(label="Detected Regions")
    
    front_button.click(
        fn=process_front_cnic,
        inputs=front_input,
        outputs=front_output
    )
    
    back_button.click(
        fn=process_back_cnic,
        inputs=back_input,
        outputs=back_output
    )


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