NVIDIA DALI从入门到放弃之三：Data Loading-CFANZ编程社区

NVIDIA DALI从入门到放弃之一：概述

NVIDIA DALI从入门到放弃之二：入门示例

NVIDIA DALI从入门到放弃之三：Data Loading

NVIDIA DALI从入门到放弃之四：Multiple GPU

NVIDIA DALI从入门到放弃之五：Image Processing

NVIDIA DALI从入门到放弃之六：Geometric Transforms

NVIDIA DALI从入门到放弃之七：Sequence Processing

NVIDIA DALI从入门到放弃之八：PyTorch Plugin API

1 自定义数据格式

import numpy as np
from random import shuffle
from nvidia.dali.pipeline import Pipeline
import nvidia.dali.fn as fn
import matplotlib.pyplot as plt
import cupy as cp
import imageio

root_path = '/home/lmin/data/aDALI/'
batch_size = 4

class CustomizeInputGpuIterator(object):
    def __init__(self, images_dir, batch_size):
        self.images_dir = images_dir
        self.batch_size = batch_size
        with open(self.images_dir + "img_list.txt", 'r') as f:
            self.files = [line.rstrip() for line in f if line != '']
        shuffle(self.files)

    def __iter__(self):
        self.idx = 0
        self.length = len(self.files)
        return self

    def __next__(self):
        batch = []
        labels = []
        for _ in range(self.batch_size):
            jpeg_filename, label = self.files[self.idx].split(' ')
            im = imageio.imread(self.images_dir + jpeg_filename)
            im = cp.asarray(im)
            im = im * 0.6
            batch.append(im.astype(cp.uint8))
            labels.append(cp.array([label], dtype = np.uint8))
            self.idx = (self.idx + 1) % self.length
        return (batch, labels)

eii_gpu = CustomizeInputGpuIterator(root_path, batch_size)
print(type(next(iter(eii_gpu))[0][0]))

pipe_gpu = Pipeline(batch_size=batch_size, num_threads=2, device_id=0)
with pipe_gpu:
    images, labels = fn.external_source(source=eii_gpu, num_outputs=2, device="gpu")
    enhance = fn.brightness_contrast(images, contrast=2)
    pipe_gpu.set_outputs(enhance, labels)

pipe_gpu.build()

pipe_out_gpu = pipe_gpu.run()
batch_gpu = pipe_out_gpu[0].as_cpu()
labels_gpu = pipe_out_gpu[1].as_cpu()

img = batch_gpu.at(2)
print(img.shape)
plt.axis('off')
plt.imshow(img)
plt.show()

2 COCO数据

from nvidia.dali.pipeline import Pipeline
import nvidia.dali.fn as fn
import nvidia.dali.types as types
import numpy as np
import os.path

file_root = '/home/lmin/data/coco/images/val2017'
annotations_file = os.path.join('/home/lmin/data/coco/annotations', 'instances_val2017.json')
batch_size = 16

pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)
with pipe:
    jpegs, bboxes, labels, polygons, vertices = fn.coco_reader(
        file_root=file_root,
        annotations_file=annotations_file,
        polygon_masks=True,
        ratio=True)
    images = fn.image_decoder(jpegs, device="mixed", output_type=types.RGB)
    pipe.set_outputs(images, bboxes, labels, polygons, vertices)


pipe.build()
pipe_out = pipe.run()

images_cpu = pipe_out[0].as_cpu()
bboxes_cpu = pipe_out[1]
labels_cpu = pipe_out[2]
polygons_cpu = pipe_out[3]
vertices_cpu = pipe_out[4]

bboxes = bboxes_cpu.at(4)
labels = labels_cpu.at(4)
for bbox, label in zip(bboxes, labels):
    x, y, width, height = bbox
    print(f"Bounding box (x={x}, y={y}, width={width}, height={height}), label={label}")

polygons = polygons_cpu.at(4)
vertices = vertices_cpu.at(4)
print(polygons.shape, vertices.shape)

for polygon in polygons:
    mask_idx, start_vertex, end_vertex = polygon
    nvertices = end_vertex - start_vertex
    print(f"Polygon belonging to mask index {mask_idx} containing {nvertices} vertices:")
    polygon_vertices = vertices[start_vertex:end_vertex]
    for vertex_id in range(nvertices):
        x, y = vertices[vertex_id]
        print(f"Vertex {vertex_id}: x={x}, y={y}")


import matplotlib.pyplot as plt
import matplotlib.patches as patches
import random

random.seed(1231243)

def plot_sample(img_index, ax):
    img = images_cpu.at(img_index)

    H = img.shape[0]
    W = img.shape[1]

    ax.imshow(img)
    bboxes = bboxes_cpu.at(img_index)
    labels = labels_cpu.at(img_index)
    polygons = polygons_cpu.at(img_index)
    vertices = vertices_cpu.at(img_index)
    categories_set = set()
    for label in labels:
        categories_set.add(label)

    category_id_to_color = dict([(cat_id , [random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1)]) for cat_id in categories_set])

    for bbox, label in zip(bboxes, labels):
        rect = patches.Rectangle((bbox[0] * W,bbox[1] * H), bbox[2] * W,bbox[3] * H,
                                 linewidth=1, edgecolor=category_id_to_color[label], facecolor='none')
        ax.add_patch(rect)

    for polygon in polygons:
        mask_idx, start_vertex, end_vertex = polygon
        polygon_vertices = vertices[start_vertex:end_vertex]
        polygon_vertices = polygon_vertices * [W, H]
        poly = patches.Polygon(polygon_vertices, True, facecolor=category_id_to_color[label], alpha=0.7)
        ax.add_patch(poly, )

fig, ax = plt.subplots(2, 2, figsize=(12, 12))
fig.tight_layout()
plot_sample(2, ax[0, 0])
plot_sample(1, ax[0, 1])
plot_sample(4, ax[1, 0])
plot_sample(8, ax[1, 1])
plt.show()