Didier/Visual_question_answering

1"""2File: app.py3 4Description: Visual question answering (with pixtral-12b-2409)5 6Author: Didier Guillevic7Date: 2024-09-298"""9import logging10logger = logging.getLogger(__name__)11 12import gradio as gr13import PIL14import base6415from io import BytesIO16 17import os18from mistralai import Mistral19 20#21# Mistral AI client22#23api_key = os.environ["MISTRAL_API_KEY"]24client = Mistral(api_key=api_key)25model_id = "pixtral-12b-2409"26 27#28# Utils29#30def image_to_base64(filepath):31    file_format = filepath.split('.')[-1].lower()32    with open(filepath, "rb") as image_file:33        base64_string = base64.b64encode(image_file.read()).decode('utf-8')34    return base64_string, file_format35 36def pil_to_base64(pil_image, image_format='jpeg'):  # Default to 'png' if format not known37    # Save the PIL image to a BytesIO object in the specified format38    buffered = BytesIO()39    pil_image.save(buffered, format=image_format)40    41    # Convert the image to a base64 string42    base64_string = base64.b64encode(buffered.getvalue()).decode('utf-8')43    44    return base64_string, image_format45 46 47#48# Generate a response49#50def answer_question(text, image, temperature=0.0, max_tokens=1_024):51    # Convert image to base64 string52    #base64_string, file_format = image_to_base64(image)53    base64_string, file_format = pil_to_base64(image)54 55    messages = [56        {57            "role": "user",58            "content": [59                {60                    "type": "text",61                    "text": text62                },63                {64                    "type": "image_url",65                    "image_url": f"data:image/{file_format};base64,{base64_string}" 66                }67            ]68        }69    ]70    71    chat_response = client.chat.complete(72        model=model_id,73        messages=messages,74        temperature=temperature,75        max_tokens=max_tokens76    )77 78    return chat_response.choices[0].message.content79 80 81def clear_all():82    # Clear question, image, output_text83    return ('', None, '')84 85 86#87# User interface88#89with gr.Blocks() as demo:90    91    gr.Markdown("""92        ## Visual question answering / image captioning93    """)94    with gr.Row():95        with gr.Column():96            question = gr.Textbox(97                placeholder="Ask anything about the image..",98                lines=2,99                render=True)100            image = gr.Image(type="pil") # return type: 'numpy', 'pil', 'filepath'101            with gr.Row():102                temperature = gr.Slider(103                    label="Temperature",104                    minimum=0.0,105                    maximum=1.0,106                    value=0.0,107                    step=0.1108                )109                max_tokens = gr.Slider(110                    label="Max tokens",111                    minimum=128,112                    maximum=2_048,113                    value=1_024,114                    step=128115                )116        with gr.Column():117            output_text = gr.Textbox(118                lines=10,119                label="Pixtral 12B",120                render=True121            )122    with gr.Row():123        clear_btn = gr.Button("Clear", variant="secondary")124        submit_btn = gr.Button(value="Submit", variant="primary")125 126    # Examples127    examples = gr.Examples(128        [129            [130                ('Can you describe this image? If this is the image a '131                 'government issued ID, extract the information and format '132                 'it as a JSON document.'),133                './sample_ID.jpeg'134            ],135            [136                'In which city and country was  this picture taken?',137                './16-park-ave-milton-park-1024x683.jpg'138            ],139            [   (140                    'Please extract the identification information present in the image. '141                    'Format the response in a JSON format such as "{name_of_element : [value]}".'142                ),143                './driver_license.png'144            ],145        ],146        inputs=[question, image],147        outputs=[output_text,],148        fn=answer_question,149        cache_examples=False,150        label="Examples"151    )152 153    # Click actions154    clear_btn.click(155        clear_all,156        inputs=[],157        outputs=[question, image, output_text]158    )159    submit_btn.click(160        fn=answer_question,161        inputs=[question, image, temperature, max_tokens],162        outputs=[output_text,]163    )164 165    # Documentation166    with gr.Accordion("Documentation", open=False):167        gr.Markdown("""168            - model: Serving pixtral-12b (https://mistral.ai/news/pixtral-12b/)169            - temperature: 0.0=little variable in output, 1.0=lots of variability in output170            - max_tokens: maximum number of tokens to generate as output (1 word approx 1.5 tokens)171        """)172 173demo.launch(show_api=False)