PaliGemma 2 Mix : guide avec projet de démo OCR

PaliGemma 2 Mix est un modèle d'IA multimodal développé par Google. Il s'agit d'une version améliorée du modèle vision-language (VLM) PaliGemma, intégrant des capacités avancées des modèles de vision SigLIP et des modèles de langage Gemma 2.

Dans ce tutoriel, je vous explique comment utiliser PaliGemma 2 Mix pour construire un scanneur de factures et un analyseur de dépenses propulsé par l'IA, capable de :

• Extraire et classer les dépenses à partir de tickets de caisse.
• Réaliser de la reconnaissance optique de caractères (OCR) pour récupérer les informations clés.
• Résumer les dépenses à partir d'images fournies.

Même si nous nous focalisons ici sur un outil d'insights financiers, vous pourrez réutiliser ce que vous apprendrez pour explorer d'autres cas d'usage de PaliGemma 2 Mix, comme la segmentation d'images, la détection d'objets ou les questions-réponses.

Qu’est-ce que PaliGemma 2 Mix ?

PaliGemma 2 Mix est un modèle vision-language (VLM) avancé qui traite à la fois des images et du texte en entrée et génère des sorties textuelles. Il est conçu pour gérer une large variété de tâches d'IA multimodale tout en prenant en charge plusieurs langues.

PaliGemma 2 est conçu pour un large éventail de tâches vision-language, notamment la légendage d'images et de courtes vidéos, les questions-réponses visuelles, l'OCR, la détection d'objets et la segmentation.

Source des images utilisées dans le schéma : Google

Le modèle PaliGemma 2 Mix est conçu pour :

• Légendage d'images et de courtes vidéos : générer des légendes précises et contextualisées pour des images statiques et de courtes vidéos.
• Questions-réponses visuelles (VQA) : analyser des images et répondre à des questions textuelles fondées sur le contenu visuel.
• Reconnaissance optique de caractères (OCR) : extraire et interpréter du texte à partir d'images, utile pour des documents, des tickets et des contenus numérisés.
• Détection et segmentation d'objets : identifier, étiqueter et segmenter les objets au sein d'une image pour une analyse structurée.
• Prise en charge multilingue : génération et compréhension de texte dans plusieurs langues pour des applications internationales.

Vous trouverez plus d'informations sur le modèle PaliGemma 2 Mix dans l'article de présentation officiel. 

Vue d'ensemble du projet : scanneur de factures et analyseur de dépenses avec PaliGemma 2 Mix

Voici les principales étapes que nous allons suivre :

1. Charger et préparer le jeu de données : commencer par charger et préparer les images de tickets en entrée.
2. Initialiser le modèle PaliGemma 2 Mix : configurer et charger le modèle pour traiter des tâches vision-language.
3. Traiter les images d'entrée : convertir les images au format adéquat (RGB) et les préparer pour l'analyse.
4. Extraire les informations clés : effectuer l'OCR pour récupérer le montant total.
5. Catégoriser les dépenses : classer les achats en catégories comme alimentation, habillement, électronique, etc.
6. Générer des insights de dépenses : résumer les dépenses catégorisées et générer un graphique de répartition.
7. Construire une interface interactive Gradio : créer une interface où les utilisateurs peuvent téléverser plusieurs factures, extraire les données et analyser visuellement les dépenses.

Étape 1 : prérequis

Avant de commencer, vérifions que les outils et bibliothèques suivants sont installés :

• Python 3.8+
• torch
• Transformers
• PIL
• Matplotlib
• Gradio

Exécutez les commandes suivantes pour installer les dépendances nécessaires :

pip install gradio -U bitsandbytes -U transformers -q

Une fois ces dépendances installées, lancez les imports suivants :

import gradio as gr
import torch
import pandas as pd
import matplotlib.pyplot as plt
from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor, BitsAndBytesConfig
from transformers import BitsAndBytesConfig
from PIL import Image
import re

Étape 2 : initialisation du modèle

Nous configurons et chargeons le modèle PaliGemma 2 Mix avec quantification pour optimiser les performances. Pour cette démo, nous utiliserons le modèle 10 milliards de paramètres avec une résolution d'entrée de 448 x 448. Vous aurez besoin au minimum d'un GPU T4 avec 40 Go de mémoire (configuration Colab) pour l'exécuter. 

device = "cuda" if torch.cuda.is_available() else "cpu"
# Model setup
model_id = "google/paligemma2-10b-mix-448" 
bnb_config = BitsAndBytesConfig(
    load_in_8bit=True,  # Change to load_in_4bit=True for even lower memory usage
    llm_int8_threshold=6.0,
)

# Load model with quantization
model = PaliGemmaForConditionalGeneration.from_pretrained(
    model_id, quantization_config=bnb_config
).eval()

# Load processor
processor = PaliGemmaProcessor.from_pretrained(model_id)
# Print success message
print("Model and processor loaded successfully!")

La quantification BitsAndBytes permet de réduire l'empreinte mémoire tout en préservant les performances, ce qui rend possible l'exécution de grands modèles sur des GPU limités. Dans cette implémentation, nous utilisons une quantification en 4 bits pour optimiser davantage l'efficacité mémoire.

Nous chargeons le modèle via la classe PaliGemmaForConditionalGeneration de la bibliothèque transformers en lui passant l'ID du modèle et la configuration de quantification. De même, nous chargeons le processor, qui prétraite les entrées en tenseurs avant de les soumettre au modèle.

Étape 3 : traitement des images 

Une fois les shards du modèle chargés, nous préparons les images avant de les passer au modèle afin d'assurer la compatibilité de format et l'homogénéité. Nous convertissons les images en RGB :

def ensure_rgb(image: Image.Image) -> Image.Image:
    if image.mode != "RGB":
        image = image.convert("RGB")
    return image

Nos images sont désormais prêtes pour l'inférence.

Étape 4 : inférence avec PaliGemma

Nous définissons à présent la fonction principale d'inférence. Elle prend des images et des questions, les intègre dans des prompts et les transmet au modèle via le processor pour l'inférence.

def ask_model(image: Image.Image, question: str) -> str:
    prompt = f"<image> answer en {question}"
    inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
    with torch.inference_mode():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=50,
            do_sample=False
        )
    result = processor.batch_decode(generated_ids, skip_special_tokens=True)
    return result[0].strip()

Étape 5 : extraction des informations clés

Maintenant que la fonction principale est prête, extrayons les paramètres clés de l'image : dans notre cas, le montant total et la catégorie de produits.

 def extract_total_amount(image: Image.Image) -> float:
    question = "ocr\n what is the total amount? in numbers only" 
    answer = ask_model(image, question)
    print(f"Answer from model: {answer}")
    try:
        amounts = re.findall(r'\d+\.\d+|\d+', answer)  # Capture both integer and decimal values
        if amounts:
            return float(amounts[-1])  # Get the last valid amount as the total
    except ValueError:
        pass
    return 0.0

La fonction extract_total_amount() traite une image pour extraire, via l'OCR, le montant total figurant sur un ticket. Elle construit une requête demandant au modèle de renvoyer uniquement des valeurs numériques, puis appelle ask_model() pour générer la réponse.

def categorize_goods(image: Image.Image) -> str:
    question = "what is the category of goods in the image - Grocery/ Clothing/ Electronics/ Other?"
    answer = ask_model(image, question)
    print(f"Category from model: {answer}")
    answer = answer.split("\n")[-1].strip().capitalize()
    valid_categories = ["Grocery", "Clothing", "Electronics", "Other"]
    return answer if answer in valid_categories else "Other"

La fonction categorize_goods() classe le type de produits présents sur une image en posant au modèle une question prédéfinie listant les catégories possibles : alimentation, habillement, électronique ou autre. ask_model() traite l'image et renvoie une réponse textuelle. Si cette réponse correspond à l'une des catégories valides, elle est renvoyée ; sinon, la catégorie "Other" est appliquée par défaut.

Étape 6 : analyse des informations

Toutes les fonctions clés sont prêtes ; analysons maintenant les résultats.

def generate_spending_chart(categories: dict):
    filtered_categories = {k: v for k, v in categories.items() if v > 0}  # Remove zero-value categories
    labels = list(filtered_categories.keys())
    values = list(filtered_categories.values())
    if not values or sum(values) == 0:
        fig, ax = plt.subplots()
        ax.text(0.5, 0.5, "No Spending Data", ha="center", va="center", fontsize=12)
        ax.axis("off")
        return fig
    fig, ax = plt.subplots()
    ax.pie(values, labels=labels, autopct='%1.1f%%', startangle=90)
    ax.axis('equal')
    plt.title("Spending Distribution")  
    return fig

Cette fonction crée un graphique en secteurs pour visualiser la répartition des dépenses par catégorie. En l'absence de données valides, elle génère une figure vierge avec le message "No Spending Data". Sinon, elle produit un camembert avec les étiquettes de catégorie et les pourcentages, pour une visualisation proportionnelle et lisible.

Étape 6 : analyser plusieurs factures simultanément

Nous avons généralement plusieurs factures à traiter ; créons une fonction pour toutes les analyser en une seule fois.

def process_multiple_bills(files: list):
    results = []
    images = []
    total_spending = 0
    category_totals = {"Grocery": 0, "Clothing": 0, "Electronics": 0, "Other": 0}
    for file in files:
        image = Image.open(file)
        image = ensure_rgb(image)
        images.append(image)
        total_amount = extract_total_amount(image)
        category = categorize_goods(image)
        total_spending += total_amount
        category_totals[category] += total_amount
        results.append({"Bill": len(results) + 1, "Category": category, "Total Amount": f"₹{total_amount:.2f}"})
    pie_chart = generate_spending_chart(category_totals)
    summary_text = f"**Total Spending Across All Bills:** ₹{total_spending:.2f}"
    return images, pd.DataFrame(results), summary_text, pie_chart

Pour analyser plusieurs factures en parallèle, nous suivons les étapes suivantes :

1. Initialiser les structures : créer des listes pour stocker les résultats et images, initialiser total_spending à 0 et définir un dictionnaire pour les totaux par catégorie.
2. Traiter chaque facture :

• Ouvrir et convertir l'image en RGB.
• Ajouter l'image à la liste.
• Extraire le montant total du ticket.
• Catégoriser les produits du ticket.
• Mettre à jour le total des dépenses et les totaux par catégorie.
• Enregistrer les données extraites dans une liste de résultats.

3. Générer des insights : créer un camembert de répartition des dépenses et un récapitulatif du total.
4. Retourner les résultats : renvoyer la liste d'images, un DataFrame de synthèse, le résumé du total des dépenses et le graphique.

Étape 7 : construire l'interface Gradio

Toutes les fonctions métier sont en place. Passons à la création d'une interface interactive avec Gradio.

def gradio_demo():
    with gr.Blocks() as demo:
        gr.Markdown("## PaliGemma 2 Mix Powered- Multiple Bill Scanner\nUpload multiple bill images, and this demo will extract text, categorize spending, and generate insights.")
        with gr.Row():
            with gr.Column():
                image_input = gr.File(file_count="multiple", file_types=["image"], label="Upload Bill Images")
                submit_button = gr.Button("Process Bills")
            with gr.Column():
                image_output = gr.Gallery(label="Uploaded Bills")
                table_output = gr.Dataframe(label="Bill Summary")
                summary_output = gr.Text(label="Total Spending Summary")
                chart_output = gr.Plot(label="Aggregated Spending Distribution")
        submit_button.click(
            fn=process_multiple_bills,
            inputs=image_input,
            outputs=[image_output, table_output, summary_output, chart_output]
        )
    demo.launch(debug=True)

if __name__ == "__main__":
    gradio_demo()

Ce code crée une UI Gradio structurée avec un téléverseur multi-images et un bouton pour lancer le traitement. Une fois soumis, les factures téléversées s'affichent en galerie, les données extraites apparaissent dans un tableau, le total est résumé en texte et un camembert de répartition est généré. 

La fonction relie les entrées utilisateur à process_multiple_bills() pour une extraction et une visualisation fluides. Enfin, demo.launch() démarre l'application Gradio pour une interaction en temps réel.

J'ai également testé cette démo avec deux factures au format image (facture d'achat Amazon) et obtenu les résultats suivants.

Remarque : les VLM ont du mal à extraire des nombres, ce qui peut parfois entraîner des résultats inexacts. Par exemple, le montant total de la deuxième facture ci-dessous est erroné. Cela peut se corriger en utilisant des modèles plus grands ou en affinant les modèles existants.

Conclusion

Dans ce tutoriel, nous avons créé un scanneur multi-factures propulsé par l'IA avec PaliGemma 2 Mix, qui nous aide à extraire et à catégoriser nos dépenses à partir de tickets. Nous avons exploité les capacités vision-language de PaliGemma 2 Mix pour l'OCR et la classification afin d'analyser nos dépenses en toute simplicité. N'hésitez pas à adapter ce tutoriel à votre propre cas d'usage.