许可证：Apache-2.0
语言：

• 越南语

基础模型：

• OpenGVLab/InternVL3-1B
• 5CD-AI/Vintern-1B-v3_5

任务类型：视觉问答

import numpy as np  
import torch  
import torchvision.transforms as T  
# from decord import VideoReader, cpu  
from PIL import Image  
from torchvision.transforms.functional import InterpolationMode  
from transformers import AutoModel, AutoTokenizer  

IMAGENET_均值 = (0.485, 0.456, 0.406)  
IMAGENET_标准差 = (0.229, 0.224, 0.225)  

def 构建预处理流程(输入尺寸):  
    均值, 标准差 = IMAGENET_均值, IMAGENET_标准差  
    预处理流程 = T.Compose([  
        T.Lambda(lambda 图像: 图像.convert('RGB') if 图像.mode != 'RGB' else 图像),  
        T.Resize((输入尺寸, 输入尺寸), interpolation=InterpolationMode.BICUBIC),  
        T.ToTensor(),  
        T.Normalize(mean=均值, std=标准差)  
    ])  
    return 预处理流程  

def 寻找最接近宽高比(原始宽高比, 目标比例集合, 原始宽度, 原始高度, 图像尺寸):  
    最小比例差 = float('inf')  
    最佳比例 = (1, 1)  
    面积 = 原始宽度 * 原始高度  
    for 比例 in 目标比例集合:  
        目标宽高比 = 比例[0] / 比例[1]  
        比例差 = abs(原始宽高比 - 目标宽高比)  
        if 比例差 < 最小比例差:  
            最小比例差 = 比例差  
            最佳比例 = 比例  
        elif 比例差 == 最小比例差:  
            if 面积 > 0.5 * 图像尺寸 * 图像尺寸 * 比例[0] * 比例[1]:  
                最佳比例 = 比例  
    return 最佳比例  

def 动态图像预处理(图像, 最小切片数=1, 最大切片数=12, 图像尺寸=448, 使用缩略图=False):  
    原始宽度, 原始高度 = 图像.size  
    原始宽高比 = 原始宽度 / 原始高度  

    # 生成候选宽高比集合  
    目标比例集合 = set(  
        (i, j) for n in range(最小切片数, 最大切片数 + 1) for i in range(1, n + 1) for j in range(1, n + 1) if  
        i * j <= 最大切片数 and i * j >= 最小切片数)  
    目标比例集合 = sorted(目标比例集合, key=lambda x: x[0] * x[1])  

    # 寻找最接近的宽高比  
    目标比例 = 寻找最接近宽高比(原始宽高比, 目标比例集合, 原始宽度, 原始高度, 图像尺寸)  

    # 计算目标尺寸  
    目标宽度 = 图像尺寸 * 目标比例[0]  
    目标高度 = 图像尺寸 * 目标比例[1]  
    切片数量 = 目标比例[0] * 目标比例[1]  

    # 调整图像尺寸  
    调整后图像 = 图像.resize((目标宽度, 目标高度))  
    处理后图像列表 = []  
    for i in range(切片数量):  
        裁剪框 = (  
            (i % (目标宽度 // 图像尺寸)) * 图像尺寸,  
            (i // (目标宽度 // 图像尺寸)) * 图像尺寸,  
            ((i % (目标宽度 // 图像尺寸)) + 1) * 图像尺寸,  
            ((i // (目标宽度 // 图像尺寸)) + 1) * 图像尺寸  
        )  
        # 分割图像  
        切片图像 = 调整后图像.crop(裁剪框)  
        处理后图像列表.append(切片图像)  
    assert len(处理后图像列表) == 切片数量  
    if 使用缩略图 and len(处理后图像列表) != 1:  
        缩略图 = 图像.resize((图像尺寸, 图像尺寸))  
        处理后图像列表.append(缩略图)  
    return 处理后图像列表  

def 加载图像(图像文件, 输入尺寸=448, 最大切片数=12):  
    图像 = Image.open(图像文件).convert('RGB')  
    预处理流程 = 构建预处理流程(输入尺寸=输入尺寸)  
    图像列表 = 动态图像预处理(图像, 图像尺寸=输入尺寸, 使用缩略图=True, 最大切片数=最大切片数)  
    像素值张量 = [预处理流程(图像) for 图像 in 图像列表]  
    像素值张量 = torch.stack(像素值张量)  
    return 像素值张量  

模型 = AutoModel.from_pretrained(  
    "TienAnh/Finetune_VQA_1B",  
    torch_dtype=torch.bfloat16,  
    low_cpu_mem_usage=True,  
    trust_remote_code=True,  
    use_flash_attn=False,  
).eval().cuda()  

分词器 = AutoTokenizer.from_pretrained("TienAnh/Finetune_VQA_1B", trust_remote_code=True, use_fast=False)  

测试图像 = 'test-image.jpg'  

像素值 = 加载图像(测试图像, 最大切片数=6).to(torch.bfloat16).cuda()  
生成配置 = dict(最大新标记数=1024, 是否采样=False, 束搜索数=3, 重复惩罚系数=2.5)  

问题 = '<image>\n提取图像中的关键信息并以markdown格式返回。'  

回复, 历史记录 = 模型.对话(分词器, 像素值, 问题, 生成配置, 历史记录=None, 返回历史=True)  
print(f'用户: {问题}\n助手: {回复}')  

#问题 = "其他问题......"  
#回复, 历史记录 = 模型.对话(分词器, 像素值, 问题, 生成配置, 历史记录=历史记录, 返回历史=True)  
#print(f'用户: {问题}\n助手: {回复}')