1. 外部数据源
采用外部数据源作为 Pipeline 的输入.
导入相关库:
import types
import collections
import numpy as np
from random import shuffle
from nvidia.dali.pipeline import Pipeline
import nvidia.dali.fn as fn
import nvidia.dali.types as types
batch_size = 16
1.1. 定义数据源
示例:采用无限迭代器作为数据源,如:
class ExternalInputIterator(object):
def __init__(self, batch_size):
self.images_dir = "../../data/images/"
self.batch_size = batch_size
with open(self.images_dir + "file_list.txt", 'r') as f:
self.files = [line.rstrip() for line in f if line is not '']
shuffle(self.files)
def __iter__(self):
self.i = 0
self.n = len(self.files)
return self
def __next__(self):
batch = []
labels = []
for _ in range(self.batch_size):
jpeg_filename, label = self.files[self.i].split(' ')
f = open(self.images_dir + jpeg_filename, 'rb')
batch.append(np.frombuffer(f.read(), dtype = np.uint8))
labels.append(np.array([label], dtype = np.uint8))
self.i = (self.i + 1) % self.n
return (batch, labels)
1.2. 使用 Pipeline
eii = ExternalInputIterator(batch_size)
#定义Pipeline
pipe = Pipeline(batch_size=batch_size, num_threads=2, device_id=0)
with pipe:
#管道通过调用 source 或 next(source) 来保持管道运行
jpegs, labels = fn.external_source(source=eii, num_outputs=2)
decode = fn.decoders.image(jpegs, device="mixed", output_type=types.RGB)
enhance = fn.brightness_contrast(decode, contrast=2)
pipe.set_outputs(enhance, labels)
#运行 Pipeline
pipe.build()
pipe_out = pipe.run()
#
batch_cpu = pipe_out[0].as_cpu()
labels_cpu = pipe_out[1] #labels 数据仍保存在内存里.
#import matplotlib.pyplot as plt
img = batch_cpu.at(2)
print(img.shape)
print(labels_cpu.at(2))
plt.axis('off')
plt.imshow(img)
plt.show()
1.3. GPU 输入数据
外部数据源操作子也可以是来自 CuPy 或其他任何支持 cuda array interface 的数据源的 GPU 数据.
示例如,
import cupy as cp
import imageio
class ExternalInputGpuIterator(object):
def __init__(self, batch_size):
self.images_dir = "../../data/images/"
self.batch_size = batch_size
with open(self.images_dir + "file_list.txt", 'r') as f:
self.files = [line.rstrip() for line in f if line is not '']
shuffle(self.files)
def __iter__(self):
self.i = 0
self.n = len(self.files)
return self
def __next__(self):
batch = []
labels = []
for _ in range(self.batch_size):
jpeg_filename, label = self.files[self.i].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.i = (self.i + 1) % self.n
return (batch, labels)
#
eii_gpu = ExternalInputGpuIterator(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)
print(labels_cpu.at(2))
plt.axis('off')
plt.imshow(img)
plt.show()
2. LMDB 数据源
DALI 能够处理 LMDB 格式保存的数据源的加载,如,Caffe 和 Caffe2 中所采用的格式.
2.1. Caffe LMDB
如,
from nvidia.dali.pipeline import Pipeline
import nvidia.dali.fn as fn
import nvidia.dali.types as types
import os.path
import numpy as np
from timeit import default_timer as timer
import matplotlib.pyplot as plt
test_data_root = os.environ['DALI_EXTRA_PATH'] #DALI 提供的测试数据
db_folder = os.path.join(test_data_root, 'db', 'lmdb')
batch_size = 9
#
pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)
with pipe:
jpegs, labels = fn.readers.caffe(path=db_folder)#caffe数据读取接口
images = fn.decoders.image(jpegs, device="mixed", output_type=types.RGB)
output = fn.crop_mirror_normalize(
images,
dtype=types.FLOAT,
crop=(224, 224),
mean=[0., 0., 0.],
std=[1., 1., 1.],
crop_pos_x=fn.random.uniform(range=(0, 1)),
crop_pos_y=fn.random.uniform(range=(0, 1)))
pipe.set_outputs(output, labels)
#
pipe.build()
pipe_out = pipe.run()
#可视化
import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
def show_images(image_batch):
columns = 3
rows = (batch_size + 1) // (columns)
fig = plt.figure(figsize = (20,(20 // columns) * rows))
gs = gridspec.GridSpec(rows, columns)
for j in range(rows*columns):
plt.subplot(gs[j])
plt.axis("off")
img_chw = image_batch.at(j)
img_hwc = np.transpose(img_chw, (1,2,0))/255.0
plt.imshow(img_hwc)
plt.show()
#
images, labels = pipe_out
show_images(images.as_cpu())
2.2. Caffe2 LMDB
如,
db_folder = os.path.join(test_data_root, 'db', 'c2lmdb')
#
pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)
with pipe:
jpegs, labels = fn.readers.caffe2(path=db_folder) #caffe2数据读取接口
images = fn.decoders.image(jpegs, device="mixed", output_type=types.RGB)
output = fn.crop_mirror_normalize(
images,
dtype=types.FLOAT,
crop=(224, 224),
mean=[0., 0., 0.],
std=[1., 1., 1.])
pipe.set_outputs(output, labels)
#
pipe.build()
pipe_out = pipe.run()
#
images, labels = pipe_out
show_images(images.as_cpu())
3. COCO 数据集
DALI 提供了 COCO Reader 用于 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
#
test_data_root = os.environ['DALI_EXTRA_PATH'] #DALI提供的测试数据.
file_root = os.path.join(test_data_root, 'db', 'coco', 'images')
annotations_file = os.path.join(test_data_root, 'db', 'coco', 'instances.json')
batch_size = 16
#
pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)
with pipe:
jpegs, bboxes, labels, polygons, vertices = fn.readers.coco(
file_root=file_root,
annotations_file=annotations_file,
polygon_masks=True,
ratio=True)
images = fn.decoders.image(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] #[x, y, width, height]
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}")
#分割mask
polygons = polygons_cpu.at(4)
vertices = vertices_cpu.at(4)
print(polygons.shape, vertices.shape)
#(1, 3) (26, 2)
#
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()
如: