🚀 电影图片字幕生成模型
该模型能够以电影描述的风格为任意照片生成描述。它基于电影海报和相关描述进行学习,为用户带来趣味体验。
🚀 快速开始
以下是使用该模型的示例代码:
from PIL import Image
import torch
from transformers import AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecoderModel
tokenizer = AutoTokenizer.from_pretrained("dumperize/movie-picture-captioning")
feature_extractor = ViTFeatureExtractor.from_pretrained("dumperize/movie-picture-captioning")
model = VisionEncoderDecoderModel.from_pretrained("dumperize/movie-picture-captioning")
max_length = 128
num_beams = 4
gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
image_path = 'path/to/image.jpg';
image = Image.open(image_path)
image = image.resize([224,224])
if image.mode != "RGB":
image = image.convert(mode="RGB")
pixel_values = feature_extractor(images=[image], return_tensors="pt").pixel_values
pixel_values = pixel_values.to(device)
output_ids = model.generate(pixel_values, **gen_kwargs)
preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
print([pred.strip() for pred in preds])
✨ 主要特性
- 能够以电影描述的风格为照片生成描述。
- 基于特定的数据集进行训练,生成具有特色的描述。
📚 详细文档
模型详情
模型描述
这是一个基于 VisionEncoderDecoderModel 的编码器 - 解码器模型。编码器使用了 Google/vit-base-patch16-224-in21k,解码器使用了 DeepPavlov/rubert-base-cased。
该模型在俄罗斯应用程序 Kinoposk 的电影海报和描述数据集上进行了优化,现在能够生成具有大片行话风格的描述。
模型来源
偏差、风险和局限性
尽管该模型使用的训练数据可以被认为是相当中立的,但该模型仍可能产生有偏差的预测。
训练详情
训练数据
我们从 2022 年 10 月所有俄语电影的开源数据中编译了一个数据集,来源为 kinopoisk。数据集中不包括描述非常短或非常长的电影,也排除了海报为空或非常小的电影。
预处理
该模型在 1x1080ti(11Gb)上训练了近 24 小时。
评估
该模型取得了以下评估结果:sacrebleu 6.84
评估指标
我们使用了 sacrebleu 指标进行评估。
📄 许可证
本模型采用 Apache-2.0 许可证。
属性 |
详情 |
模型类型 |
图像转文本模型 |
评估指标 |
BLEU、sacrebleu |
训练数据 |
从 2022 年 10 月所有俄语电影的开源数据中编译的数据集,排除了描述非常短或非常长的电影,以及海报为空或非常小的电影 |