Hello,
I am using Nvidia Dali as a dataloader for loading train images and labels. However, I want to adapt the code so that the dataloader will also retrieve an additional attribute of the image, e.g. the filename or some other attribute which belongs to the image and is stored in a textfile.
Let’s assume that I want to retrieve the filename, how can I modify the code in order to receive the filename of each image as a third output in addition to images and labels?
Please find my code snippet below:
from nvidia.dali.pipeline import Pipeline
import nvidia.dali.ops as ops
import nvidia.dali.types as types
from nvidia.dali.plugin.pytorch import DALIClassificationIterator, DALIGenericIterator, LastBatchPolicy
def mux(condition, true_case, false_case):
neg_condition = condition ^ True
return condition * true_case + neg_condition * false_case
class TrainPipeline(Pipeline):
def __init__(self, image_dir, batch_size=16, num_threads=6, device_id=0):
super(TrainPipeline, self).__init__(batch_size, num_threads, device_id, seed=12)
self.input = ops.FileReader(file_root = image_dir, random_shuffle = True, initial_fill = 1024, pad_last_batch=True, name="Reader_Train")
self.decode = ops.ImageDecoder(device = "mixed", output_type = types.RGB)
self.resize = ops.Resize(size=224,mode='not_larger',device="gpu")
self.paste = ops.Paste(paste_x=0.5, paste_y=0.5, ratio=1,min_canvas_size=IMG_SIZE, fill_value=(0, 0, 0),device="gpu")
self.cmnp = ops.CropMirrorNormalize(device="gpu",
dtype=types.FLOAT,
output_layout=types.NCHW,
mean=DATA_MEAN,
std=DATA_STD
)
def define_graph(self):
jpegs, labels = self.input()
images_base = self.decode(jpegs)
images_base = self.resize(images_base)
images_base = self.paste(images_base)
images = self.cmnp(images_base)
return (images, labels)
def get_train_iter(image_dir, batch_size, num_threads):
pipe_train = TrainPipeline(image_dir=os.path.join(image_dir,'train'), batch_size=batch_size, num_threads=num_threads)
pipe_train.build()
dali_iter_train = DALIClassificationIterator(pipe_train, reader_name="Reader_Train", auto_reset=True, last_batch_policy=LastBatchPolicy.PARTIAL)
return dali_iter_train
for batch_nr, data_batch in enumerate(data_iter[phase]):
inputs = data_batch[0]["data"].to(device)
labels = torch.squeeze(data_batch[0]["label"],dim=1).to(device).long()