语言: 中文 标签:

• 图像分类
• 图像描述生成

海报转剧情

这是一个通过电影/电视剧海报生成剧情描述的图像描述模型。它能生成不错的剧情概要，但远非完美。我们仍在持续改进模型。

Hugging Face Spaces在线演示: https://huggingface.co/spaces/deepklarity/poster2plot

模型详情

基础模型采用Vision Transformer (ViT)作为图像编码器，GPT-2作为解码器。

我们使用的具体模型如下：

• 编码器: google/vit-base-patch16-224-in21k
• 解码器: gpt2

数据集

训练使用了公开的IMDb数据集。

使用方法

在PyTorch中

import torch
import re
import requests
from PIL import Image
from transformers import AutoTokenizer, AutoFeatureExtractor, VisionEncoderDecoderModel

# 用于忽略两个及以上句点后的文本
正则表达式 = "[.]{2,}"


def 后处理(文本):
    try:
        文本 = 文本.strip()
        文本 = re.split(正则表达式, 文本)[0]
    except Exception as e:
        print(e)
        pass
    return 文本


def 预测(图像, 最大长度=64, 束搜索数=4):
    像素值 = 特征提取器(images=图像, return_tensors="pt").pixel_values
    像素值 = 像素值.to(设备)

    with torch.no_grad():
        输出序列 = 模型生成(
            像素值,
            max_length=最大长度,
            num_beams=束搜索数,
            return_dict_in_generate=True,
        ).sequences

    预测结果 = 分词器.batch_decode(输出序列, skip_special_tokens=True)
    最终预测 = 后处理(预测结果[0])

    return 最终预测


模型路径 = "deepklarity/poster2plot"
设备 = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 加载模型
模型 = VisionEncoderDecoderModel.from_pretrained(模型路径)
模型.to(设备)
print("模型加载完成")

特征提取器 = AutoFeatureExtractor.from_pretrained(模型.encoder.name_or_path)
print("特征提取器加载完成")

分词器 = AutoTokenizer.from_pretrained(模型.decoder.name_or_path, use_fast=True)
if 模型.decoder.name_or_path == "gpt2":
    分词器.pad_token = 分词器.eos_token

print("分词器加载完成")

测试链接 = "https://upload.wikimedia.org/wikipedia/en/2/26/Moana_Teaser_Poster.jpg"
with Image.open(requests.get(测试链接, stream=True).raw) as 图像:
    预测结果 = 预测(图像)

print(预测结果)