时间戳

首页 > 解决方案 > 如何处理目录中的多标签数据集以在 PyTorch 中添加图像字幕

python - 如何处理目录中的多标签数据集以在 PyTorch 中添加图像字幕

问题描述

我需要 PyTorch 方面的帮助，关于 Dataloader 和数据集有人可以帮助/指导我吗

这是我的查询：我正在尝试使用https://github.com/yunjey/pytorch-tutorial/tree/master/tutorials/03-advanced/image_captioning进行图像字幕。

在这里，他们使用了标准 COCO 数据集。

我有数据集作为 images/ 和 captions/ 目录。

例子

目录结构： 

images/T001.jpg 
images/T002.jpg 
...
...
captions/T001.txt
captions/T002.txt
....
....

以上是关系。字幕文件在每个单独的行中有“n”个字幕。

我能够创建一个自定义数据集类，因为正在返回完整的字幕文件内容。但我只希望返回一条单独的气体。

关于如何实现这一目标的任何指导/建议。

++++++++++++++++++++++++++++++++++++++++++++++++++++这里是我设计的课程：

from __future__ import print_function
import torch
from torchvision import datasets, models, transforms
from torchvision import transforms
from torch.autograd import Variable
from torch.nn.utils.rnn import pack_padded_sequence
import torch.optim as optim
import torch.nn as nn
#from torch import np
import numpy as np
import utils_c
from data_loader_c import get_cust_data_loader 
from models import CNN, RNN
from vocab_custom import Vocabulary, load_vocab
import os

class ImageCaptionDataSet(data.Dataset):
    def __init__(self, path, json, vocab=None, transform=None):
        self.vocab = vocab 
        self.transform = transform
        self.img_dir_path = path  
        self.cap_dir_path = json 
        self.all_imgs_path = glob.glob(os.path.join(self.img_dir_path,'*.jpg'))
        self.all_caps_path = glob.glob(os.path.join(self.cap_dir_path,'*.txt'))
        pass

    def __getitem__(self,index):
        vocab = self.vocab

        img_path = self.all_imgs_path[index]
        img_base_name = os.path.basename(img_path) 
        cap_base_name = img_base_name.replace(".jpg",".txt")
        cap_path  = os.path.join(self.cap_dir_path,cap_base_name)

        caption_all_for_a_image = open(cap_path).read().split("\n")

        image = Image.open(img_path)
        image = image.convert('RGB')

        if self.transform != None:
            # apply image preprocessing
            image = self.transform(image)

        #captions_combined = []
        #max_len = 0  
        #for caption in caption_all_for_a_image:
        #    caption_str = str(caption).lower()
        #    tokens = nltk.tokenize.word_tokenize(caption_str)
        #    m = len(tokens) + 2 
        #    if m>max_len:
        #        max_len = m 
        #    caption = torch.Tensor([vocab(vocab.start_token())] +
        #                           [vocab(token) for token in tokens] +
        #                           [vocab(vocab.end_token())])
        #    captions_combined.append(caption) 
        #    #yield image, caption
        #return image,torch.Tensor(captions_combined)

        caption_str = str(caption_all_for_a_image).lower()
        tokens = nltk.tokenize.word_tokenize(caption_str)
        caption = torch.Tensor([vocab(vocab.start_token())] +
                                   [vocab(token) for token in tokens] +
                                   [vocab(vocab.end_token())])

        return image,caption

    def __len__(self):
        return len(self.all_imgs_path)

+++++++++++++++++++++++++++++++++++

标签： pythonpytorch

解决方案

首先，使用str()将字幕列表转换为单个字符串 ( caption_str = str(caption_all_for_a_image)) 是一个坏主意：

cap = ['a sentence', 'bla bla bla']
str(cap)

返回此刺：

"['a sentence', 'bla bla bla']"

请注意['， 和', '是结果字符串的一部分！

您可以随机选择其中一个字幕：

import random
...
cap_idx = random.randi(0, len(caption_all_for_a_image)-1)  # pick one at random
caption_str = caption_all_for_a_image[cap_idx].lower()  # actual selection

推荐阅读