#In-built libraries
import json
import tempfile
import traceback
from typing import Dict

#third-party libraries
import gradio as gr
from PIL import Image
from qwen_vl_utils import process_vision_info

from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor

def save_temp_image(image: Image.Image) -> str:
    """
    Saves the given PIL Image object as a temporary PNG file.
    Args:
        image (Image.Image): The image to be saved.
    Returns:
        str: The file path of the saved temporary image.
    """
    # Create a temp file WITHOUT extension
    with tempfile.NamedTemporaryFile(suffix=".tmp", delete=False) as tmp_file:
        # Save image as PNG regardless of original format
        image.save(tmp_file.name, format="PNG")
        return tmp_file.name

def id_extractor(image: Image.Image) -> Dict:

 
    # default: Load the model on the available device(s)

    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(

        "Qwen/Qwen2.5-VL-7B-Instruct", torch_dtype="auto", device_map="auto"

    )

    # default processer

    processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")

    messages = [

        {

            "role": "user",

            "content": [

                {

                    "type": "image",

                    "image": image,

                },

                {"type": "text", "text": "Extract all the available key details from the image in JSON"},

            ],

        }

    ]

    # Preparation for inference

    text = processor.apply_chat_template(

        messages, tokenize=False, add_generation_prompt=True

    )

    image_inputs, video_inputs = process_vision_info(messages)

    inputs = processor(

        text=[text],

        images=image_inputs,

        videos=video_inputs,

        padding=True,

        return_tensors="pt",

    )


    # Inference: Generation of the output

    generated_ids = model.generate(**inputs, max_new_tokens=128)

    generated_ids_trimmed = [

        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)

    ]

    output_text = processor.batch_decode(

        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False

    )
    resp = output_text[-1].replace("```json", "").replace("```", "").strip()
    return json.loads(resp)
 
# Define the Gradio interface for the ID extractor
id_interface = gr.Interface(
    fn=id_extractor,
    inputs=gr.Image(type="pil", label="Upload an image"),
    outputs=gr.JSON(label="Extracted Details"),
    title="Upload your ID",
    description="Upload an image of a document. Key details will be extracted automatically."
)

# Launch the Gradio interface
id_interface.launch(mcp_server=True)