怎么用python生成webdataset格式的数据

2024年2月21日 10:57 by wst

数据处理

---

       为了使用大量图片进行模型训练，需要准备数据格式为webdataset格式，在代码操作的过程中总是会遇到这样或那样的问题，这里给出一些实例来帮助解决。

1通过PIL读取图片并压缩

from PIL import Image
from webdataset import TarWriter
from pathlib import Path

# 从磁盘读取图片
image_path = "/home/wst/dataset/tmp/20210106-152753_254_0__16115_1609918209412_0-1.jpg"
image = Image.open(image_path)

# 将图像转换为二进制格式
binary_data = image

# 如果PIL.Image格式要转换为cv2格式，用如下代码
# image = cv2.cvtColor(numpy.asarray(image),cv2.COLOR_RGB2BGR) 

# 把图片写入压缩包
writer = TarWriter("test.tar")
to_json = {
            "__key__": Path(image_path).stem,
            "jpg": binary_data
        }
writer.write(to_json)
writer.close()

2通过cv2读取图片并压缩

import cv2
from webdataset import TarWriter
from pathlib import Path

# 从磁盘读取图片
image_path = "/home/wst/dataset/tmp/20210106-152753_254_0__16115_1609918209412_0-1.jpg"
image = cv2.imread(image_path)

# 将图像转换为PIL.Image格式，因为webdataset中的图片默认是这个格式。
binary_data = cv2.imencode('.jpg', image)[1].tobytes()

# 把图片写入压缩包
writer = TarWriter("test.tar")
to_json = {
            "__key__": Path(image_path).stem,
            "jpg": binary_data
        }
writer.write(to_json)
writer.close()

总结

通过阅读源码发现，在图片写入过程中是有校验的：

import PIL
assert isinstance(image, (PIL.Image.Image, np.ndarray)), type(image)

即 图片需要符合两种格式中的任一种：PIL.Image.Image 或者 np.ndarray

---