app.py

app.py

@@ -0,0 +1,80 @@
+#In-built libraries
+import json
+import tempfile
+import traceback
+from typing import Dict
+
+#third-party libraries
+import gradio as gr
+import ollama
+from PIL import Image
+
+
+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:
+    """
+    Extracts key details from the provided image using the ollama chat model.
+
+    Args:
+        image (Image.Image): The image from which to extract details.
+
+    Returns:
+        Dict: A dictionary containing the extracted details.
+        If the image is None or an error occurs, returns an empty dictionary.
+    """
+    try:
+        error_trace = None
+
+        if image is None:
+            # Return empty dictionary and make the output invisible
+            return {}, gr.update(visible=False)
+        
+        # Save the image temporarily
+        image_path = save_temp_image(image)
+
+        # Send the image to the ollama chat model for processing
+        response = ollama.chat(
+            model='qwen2.5vl:7b',
+            messages=[{
+                'role': 'user',
+                'content': "Extract key details like 'name', 'date of birth', 'ID number', 'Issuer' from the image as JSON, excluding signatures.",
+                'images': [image_path]
+            }]
+        )
+        
+        # Clean up the response content
+        resp = response.message.content.replace("```json", "").replace("```", "").strip()
+        return json.loads(resp)
+    
+    except json.JSONDecodeError as e:
+        # Capture and print the error traceback
+        error_trace = traceback.format_exc()
+        print(error_trace)
+        return "Kindly upload an image with good clarity"
+
+# 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)