Spaces:
Running
Running
| import gradio as gr | |
| from transformers import ( | |
| BlipProcessor, | |
| BlipForConditionalGeneration, | |
| BlipForQuestionAnswering | |
| ) | |
| print("Loading models...") | |
| caption_processor = BlipProcessor.from_pretrained( | |
| "Salesforce/blip-image-captioning-base" | |
| ) | |
| caption_model = BlipForConditionalGeneration.from_pretrained( | |
| "Salesforce/blip-image-captioning-base" | |
| ) | |
| vqa_processor = BlipProcessor.from_pretrained( | |
| "Salesforce/blip-vqa-base" | |
| ) | |
| vqa_model = BlipForQuestionAnswering.from_pretrained( | |
| "Salesforce/blip-vqa-base" | |
| ) | |
| print("Models Loaded!") | |
| def analyze_image(image, question): | |
| if image is None: | |
| return "", "" | |
| caption_inputs = caption_processor( | |
| image, | |
| return_tensors="pt" | |
| ) | |
| caption_ids = caption_model.generate( | |
| **caption_inputs, | |
| max_new_tokens=30 | |
| ) | |
| caption = caption_processor.decode( | |
| caption_ids[0], | |
| skip_special_tokens=True | |
| ) | |
| answer = "" | |
| if question and question.strip(): | |
| vqa_inputs = vqa_processor( | |
| image, | |
| question, | |
| return_tensors="pt" | |
| ) | |
| answer_ids = vqa_model.generate( | |
| **vqa_inputs, | |
| max_new_tokens=10 | |
| ) | |
| answer = vqa_processor.decode( | |
| answer_ids[0], | |
| skip_special_tokens=True | |
| ) | |
| return caption, answer | |
| with gr.Blocks( | |
| title="Visual AI Assistant" | |
| ) as demo: | |
| gr.Markdown( | |
| """ | |
| # Visual AI Assistant | |
| This application uses BLIP models for: | |
| - Image Captioning | |
| - Visual Question Answering (VQA) | |
| Upload an image and optionally ask a question about it. | |
| """ | |
| ) | |
| image_input = gr.Image( | |
| type="pil", | |
| label="Upload Image" | |
| ) | |
| question_input = gr.Textbox( | |
| label="Ask a Question", | |
| placeholder="What is happening in this image?" | |
| ) | |
| with gr.Row(): | |
| analyze_btn = gr.Button( | |
| "Analyze Image" | |
| ) | |
| clear_btn = gr.ClearButton( | |
| [ | |
| image_input, | |
| question_input | |
| ] | |
| ) | |
| caption_output = gr.Textbox( | |
| label="Generated Caption" | |
| ) | |
| answer_output = gr.Textbox( | |
| label="Answer" | |
| ) | |
| analyze_btn.click( | |
| fn=analyze_image, | |
| inputs=[ | |
| image_input, | |
| question_input | |
| ], | |
| outputs=[ | |
| caption_output, | |
| answer_output | |
| ] | |
| ) | |
| demo.launch() |