49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
|
import torch
|
|||
|
|
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
|
|||
|
|
from PIL import Image
|
|||
|
|
import os, sys
|
|||
|
|
|
|||
|
|
# 加载模型和处理器（替换为你自己的模型路径）
|
|||
|
|
model_path = "./qwen2-vl-2b"
|
|||
|
|
|
|||
|
|
model = Qwen2VLForConditionalGeneration.from_pretrained(
|
|||
|
|
model_path,
|
|||
|
|
device_map="auto", # 自动分配设备（CPU/GPU）
|
|||
|
|
torch_dtype=torch.float16 # 使用半精度加快推理
|
|||
|
|
)
|
|||
|
|
processor = AutoProcessor.from_pretrained(model_path)
|
|||
|
|
|
|||
|
|
# 加载本地图片（替换成你自己的图片路径）
|
|||
|
|
image_path = "./example.jpg"
|
|||
|
|
if not os.path.exists(image_path):
|
|||
|
|
raise FileNotFoundError(f"Image file {image_path} not found!")
|
|||
|
|
|
|||
|
|
image = Image.open(image_path).convert("RGB")
|
|||
|
|
|
|||
|
|
# 示例任务1：图像描述生成
|
|||
|
|
prompt = "Describe the content of this image in detail."
|
|||
|
|
|
|||
|
|
# 构建输入
|
|||
|
|
inputs = processor(text=prompt, images=image, return_tensors="pt").to(model.device)
|
|||
|
|
|
|||
|
|
# 推理
|
|||
|
|
with torch.no_grad():
|
|||
|
|
generated_ids = model.generate(**inputs, max_new_tokens=100)
|
|||
|
|
|
|||
|
|
# 解码输出
|
|||
|
|
description = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
|
|||
|
|
print("【图像描述】:", description)
|
|||
|
|
|
|||
|
|
# 示例任务2：视觉问答（VQA）
|
|||
|
|
question = "What is the main object in the image?"
|
|||
|
|
|
|||
|
|
# 构建输入
|
|||
|
|
inputs = processor(text=question, images=image, return_tensors="pt").to(model.device)
|
|||
|
|
|
|||
|
|
# 推理
|
|||
|
|
with torch.no_grad():
|
|||
|
|
generated_ids = model.generate(**inputs, max_new_tokens=50)
|
|||
|
|
|
|||
|
|
# 解码输出
|
|||
|
|
answer = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
|
|||
|
|
print("【问题回答】:", answer)
|