Spaces:
Sleeping
Sleeping
File size: 2,950 Bytes
25af3cf 072674d 25af3cf 072674d 25af3cf 072674d 25af3cf 072674d 25af3cf 072674d 25af3cf 072674d 25af3cf 072674d 25af3cf 072674d 25af3cf 072674d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | import base64
import gradio as gr
import io
import os
from PIL import Image, ImageEnhance, ImageFilter
from mistralai.client import Mistral
from dotenv import load_dotenv
load_dotenv()
client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])
def preprocess(image: Image.Image) -> Image.Image:
# Upscale if the image is small (photos taken from distance / low-res)
w, h = image.size
if w < 1500:
image = image.resize((w * 2, h * 2), Image.LANCZOS)
# Convert to grayscale — removes colour noise irrelevant to text
image = image.convert("L")
# Sharpen edges before contrast so fine strokes are preserved
image = image.filter(ImageFilter.SHARPEN)
# Boost contrast to make text pop against background
image = ImageEnhance.Contrast(image).enhance(2.0)
# Second sharpness pass to crisp up letter edges
image = ImageEnhance.Sharpness(image).enhance(2.0)
return image
def image_to_base64(image: Image.Image) -> str:
buffer = io.BytesIO()
image.save(buffer, format="JPEG", quality=95)
return base64.b64encode(buffer.getvalue()).decode("utf-8")
def extract_info(image: Image.Image) -> str:
image = preprocess(image)
b64 = image_to_base64(image)
response = client.chat.complete(
model="pixtral-12b-2409",
messages=[
{
"role": "system",
"content": (
"Tu es un assistant d'extraction de données. "
"Extrait les informations demandées et renvoie un objet JSON propre. "
"N'inclue aucune explication ni mise en forme Markdown. "
"Formate le résultat en JSON simple en utilisant les clés suivantes uniquement: "
"Nom, Prénom, Numéro de voie, Type de voie, Nom de la voie, "
"Complément du numéro de voie, Complément d'adresse, Code postal, Ville. "
"N'incorpore aucune donnée supplémentaire. "
"Si tu ne trouves pas la donnée, indique n/a en valeur."
),
},
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": f"data:image/jpeg;base64,{b64}",
},
{
"type": "text",
"text": "Extrait les informations structurées de ce document et renvoie uniquement le JSON.",
},
],
},
],
)
return response.choices[0].message.content
demo = gr.Interface(
fn=extract_info,
inputs=gr.Image(type="pil", label="Upload Image"),
outputs=gr.Textbox(label="Extracted JSON"),
title="Doc Xtract",
description="Upload an image to extract structured information as JSON.",
)
if __name__ == "__main__":
demo.launch(server_name="0.0.0.0")
|