高效的圖像加載

在編寫優化代碼時,圖像加載在計算機視覺中起着重要作用。此過程可能是許多CV任務的瓶頸,並且通常可能是導致性能不佳的罪魁禍首。我們需要儘可能快地從磁盤獲取圖像。

這項任務的重要性最明顯的例子是在任何CNN訓練框架中實現Dataloader類。快速加載圖像至關重要。如果不是這樣,訓練過程將受到CPU的限制,並浪費寶貴的GPU時間。

今天,我們將看一些Python庫,這些庫使我們能夠最高效地讀取圖像。他們是:

  • OpenCV
  • Pillow
  • Pillow-SIMD
  • TurboJpeg

此外,我們還將介紹使用以下從數據庫加載圖像的替代方法:

  • LMDB
  • TFRecords

最後,我們將比較每個圖片的加載時間,並找出哪個是獲勝者!

1 安裝

在開始之前,我們需要創建一個虛擬環境:

$ virtualenv -p python3.7 venv
$ source venv/bin/activate

然後,安裝所需的庫:

$ pip install -r requirements.txt

2 加載圖像的方式

2.1 Structure

通常,我們需要加載存儲在數據庫中或作爲文件夾存儲的多個圖像。在我們的場景中,抽象圖像加載器應該能夠存儲指向這樣的數據庫或文件夾的路徑,並一次從中加載一個圖像。此外,我們需要測量代碼某些部分的時間。(可選)在加載開始之前可能需要進行一些初始化。我們的ImageLoader類如下所示:

import os
from abc import abstractmethod


class ImageLoader:
    extensions: tuple = \
        (".png", ".jpg", ".jpeg", ".tiff", ".bmp", ".gif", ".tfrecords")

    def __init__(self, path: str, mode: str = "BGR"):
        self.path = path
        self.mode = mode
        self.dataset = self.parse_input(self.path)
        self.sample_idx = 0

    def parse_input(self, path):

        # single image or tfrecords file
        if os.path.isfile(path):
            assert path.lower().endswith(
                self.extensions,
            ), f"Unsupportable extension, please, use one of 
                 {self.extensions}"
            return [path]

        if os.path.isdir(path):
            # lmdb environment
            if any([file.endswith(".mdb") for file in os.listdir(path)]):
                return path
            else:
                # folder with images
                paths = \
                    [os.path.join(path, image) for image in os.listdir(path)]
                return paths

    def __iter__(self):
        self.sample_idx = 0
        return self

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

    @abstractmethod
    def __next__(self):
        pass

不同庫中的圖像解碼函數可以返回不同格式的圖像-RGB或BGR。在我們的情況下,我們默認使用BGR色彩模式,但始終可以將其轉換爲所需的格式。如果您想知道OpenCV使用BGR格式的有趣原因,請單擊此鏈接

現在,我們可以從基類繼承新類,並將其用於我們的任務。

2.2 OpenCV

第一個是OpenCV庫。我們可以使用一個簡單的函數從磁盤讀取圖像——cv2.imread

import cv2

class CV2Loader(ImageLoader):
    def __next__(self):
        start = timer()
        # get image path by index from the dataset
        path = self.dataset[self.sample_idx]
        # read the image 
        image = cv2.imread(path)
        full_time = timer() - start

        if self.mode == "RGB":
            start = timer()
            # change color mode
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  
            full_time += timer() - start

        self.sample_idx += 1
        return image, full_time

在圖像可視化之前,我們需要提到OpenCV的cv2.imshow函數需要BGR格式的圖像。一些庫默認使用RGB圖像模式,在這種情況下,我們將圖像轉換爲BGR以實現正確的可視化。

要測試OpenCV庫,請使用以下命令:

python3 show_image.py --path images/cat.jpg --method cv2

文本中的此命令和下一個命令將使用不同的庫向您顯示圖像及其加載時間。

如果一切順利,您將在窗口中看到如下圖像:

在這裏插入圖片描述

另外,您可以顯示文件夾中的所有圖像。除了使用特定的圖像,還可以使用包含圖像的文件夾的路徑:

$ python3 show_image.py --path images/pexels --method cv2

這將一次顯示文件夾中的所有圖像及其加載時間。要停止演示,可以按ESC按鈕。

2.3 Pillow

現在讓我們嘗試一下PIL庫。我們可以使用Image.open函數讀取圖像。

import numpy as np
from PIL import Image

class PILLoader(ImageLoader):
    def __next__(self):
        start = timer()
        # get image path by index from the dataset     
        path = self.dataset[self.sample_idx]  
        # read the image as numpy array
        image = np.asarray(Image.open(path))  
        full_time = timer() - start

        if self.mode == "BGR":
            start = timer()
            # change color mode
            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  
            full_time += timer() - start

        self.sample_idx += 1
        return image, full_time

我們還將Image對象轉換爲Numpy數組,因爲下一步可能需要應用一些增強或預處理,因此Numpy是其默認選擇。

要在單個圖像上進行檢查,可以使用:

$ python3 show_image.py --path images/cat.jpg --method pil

如果要在帶有圖像的文件夾中使用它:

$ python3 show_image.py --path images/pexels --method pil

2.4 Pillow-SIMD

Pillow庫的 fork-follower 具有更高的性能。Pillow-SIMD使用新技術,可以使用與標準Pillow相同的API更快地讀取和轉換圖像。

不能在同一虛擬環境中同時使用Pillow和Pillow-SIMD,默認情況下將使用Pillow-SIMD。

要使用Pillow-SIMD並避免由於Pillow和Pillow-SIMD在一起而導致的錯誤,您需要創建一個新的虛擬環境並使用:

$ pip install pillow-simd

或者,您可以卸載以前的Pillow版本並安裝Pillow-SIMD:

$ pip uninstall pillow
$ pip install pillow-simd

您無需更改代碼中的任何內容,前面的示例仍然有效。要檢查一切是否正常,可以使用上一個“Pillow”部分中的命令:

$ python3 show_image.py --path images/cat.jpg --method pil
$ python3 show_image.py --path images/pexels --method pil

2.5 TurboJpeg

還有另一個名爲TurboJpeg的庫。正如標題所示,它只能讀取使用JPEG壓縮的圖像。

讓我們使用TurboJpeg創建一個圖像加載器。

from turbojpeg import TurboJPEG

class TurboJpegLoader(ImageLoader):
    def __init__(self, path, **kwargs):
        super(TurboJpegLoader, self).__init__(path, **kwargs)
        # create TurboJPEG object for image reading
        self.jpeg_reader = TurboJPEG()  

    def __next__(self):
        start = timer()
        # open the input file as bytes
        file = open(self.dataset[self.sample_idx], "rb")  
        full_time = timer() - start

        if self.mode == "RGB":
            mode = 0
        elif self.mode == "BGR":
            mode = 1

        start = timer()
        # decode raw image
        image = self.jpeg_reader.decode(file.read(), mode)  
        full_time += timer() - start

        self.sample_idx += 1
        return image, full_time

TurboJpeg需要對輸入圖像進行解碼,並將其存儲爲字節字符串。

您可以使用以下命令嘗試。但是請記住,TurboJpeg僅允許處理.jpeg圖像:

$ python3 show_image.py --path images/cat.jpg --method turbojpeg
$ python3 show_image.py --path images/pexels --method turbojpeg

2.6 LMDB

當優先考慮速度時,通常使用的圖像加載方法是事先將數據轉換爲更好的表示形式(數據庫或序列化緩衝區)。這種“數據庫”的最大優勢之一是,它們每次數據訪問的系統調用數爲零,而文件系統每次數據訪問則需要多個系統調用。我們可以創建一個LMDB數據庫,該數據庫將以鍵值格式收集所有圖像。

以下函數使我們可以使用圖像創建LMDB環境。LMDB的“環境”實質上是一個文件夾,其中包含由LMDB庫創建的特殊文件。此函數只需要包含圖像路徑和保存路徑的列表:

import cv2
import lmdb
import numpy as np

def store_many_lmdb(images_list, save_path):
    
    # number of images in our folder
    num_images = len(images_list)  
    # all file sizes
    file_sizes = [os.path.getsize(item) for item in images_list]  
    # the maximum file size index
    max_size_index = np.argmax(file_sizes)  
    # maximum database size in bytes
    map_size = num_images * cv2.imread(images_list[max_size_index]).nbytes * 10

    # create lmdb environment
    env = lmdb.open(save_path, map_size=map_size)  

    # start writing to environment
    with env.begin(write=True) as txn:  
        for i, image in enumerate(images_list):
            with open(image, "rb") as file:
                # read image as bytes
                data = file.read()  
                 # get image key
                key = f"{i:08}" 
                # put the key-value into database
                txn.put(key.encode("ascii"), data)  

    # close the environment
    env.close()  

有一個Python腳本可使用圖像創建LMDB環境:

  • --path參數應包含您收集的圖像文件夾的路徑
  • --output參數是將在其中創建LMDB的目錄
$ python3 create_lmdb.py --path images/pexels --output lmdb/images

現在,隨着LMDB環境的創建,我們可以從中加載圖像。讓我們創建一個新的加載程序類。

在從數據庫加載圖像的情況下,我們需要打開該數據庫進行讀取。有一個名爲open_database的新函數。它返回迭代器以瀏覽打開的數據庫。另外,當此迭代器到達數據末尾時,我們需要使用_iter_函數將其返回到數據庫的開頭。

LMDB允許我們存儲數據,但是沒有內置的圖像解碼器。由於缺少解碼器,我們將在此處使用cv2.imdecode函數。

class LmdbLoader(ImageLoader):
    def __init__(self, path, **kwargs):
        super(LmdbLoader, self).__init__(path, **kwargs)
        self.path = path
        self._dataset_size = 0
        self.dataset = self.open_database()

    # we need to open the database to read images from it
    def open_database(self):
        # open the environment by path
        lmdb_env = lmdb.open(self.path)  
        # start reading
        lmdb_txn = lmdb_env.begin()  
        # create cursor to iterate through the database
        lmdb_cursor = lmdb_txn.cursor() 
        # get number of items in full dataset
        self._dataset_size = lmdb_env.stat()["entries"]  
        return lmdb_cursor

    def __iter__(self):
        # set the cursor to the first database element
        self.dataset.first()  
        return self

    def __next__(self):
        start = timer()
        # get raw image
        raw_image = self.dataset.value()  
        # convert it to numpy
        image = np.frombuffer(raw_image, dtype=np.uint8)  
        # decode image
        image = cv2.imdecode(image, cv2.IMREAD_COLOR)  
        full_time = timer() - start

        if self.mode == "RGB":
            start = timer()
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            full_time += timer() - start

        start = timer()
        # step to the next element in database
        self.dataset.next()  
        full_time += timer() - start
        return image, full_time

    def __len__(self):
        # get dataset length
        return self._dataset_size

創建環境和加載程序類之後,我們可以檢查其正確性並顯示其中的圖像。現在,在--path參數中,我們需要提及LMDB環境的路徑。請記住,您可以使用ESC按鈕停止顯示。

$ python3 show_image.py --path lmdb/images --method lmdb

2.7 TFRecords

另一個有用的數據庫是TFRecords。爲了有效地讀取數據,對數據進行序列化並將其存儲在一組可以線性讀取的文件(每個100-200MB)中會很有幫助(TensorFlow手冊)。

在創建tfrecords文件之前,我們需要選擇數據庫的結構。TFRecords允許保留具有許多其他功能的項目。如果需要,可以保存文件名或圖像的寬度和高度。所有這些東西都應該收集在python字典中,即:

image_feature_description = {
    "height" :tf.io.FixedLenFeature([], tf.int64),
    "width" :tf.io.FixedLenFeature([], tf.int64),
    "filename": tf.io.FixedLenFeature([], tf.string),
    "label": tf.io.FixedLenFeature([], tf.int64),
    "image_raw": tf.io.FixedLenFeature([], tf.string),
}

在我們的示例中,我們將僅使用原始字節格式的圖像及其唯一的鍵,稱爲“標籤”。

import os
import tensorflow as tf

def _byte_feature(value):
    """Convert string / byte into bytes_list."""
    if isinstance(value, type(tf.constant(0))):
        # BytesList can't unpack string from EagerTensor.
        value = value.numpy() 
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def _int64_feature(value):
    """Convert bool / enum / int / uint into int64_list."""
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

def image_example(image_string, label):
    feature = {
        "label": _int64_feature(label),
        "image_raw": _byte_feature(image_string),
    }
    return tf.train.Example(features=tf.train.Features(feature=feature))

def store_many_tfrecords(images_list, save_file):

    assert save_file.endswith(
        ".tfrecords"
    ), 'File path is wrong, it should contain "*myname*.tfrecords"'

    directory = os.path.dirname(save_file)
    if not os.path.exists(directory):
        os.makedirs(directory)
    # start writer
    with tf.io.TFRecordWriter(save_file) as writer: 
        # cycle by each image path
        for label, filename in enumerate(images_list): 
            # read the image as bytes string
            image_string = open(filename, "rb").read()  
            # save the data as tf.Example object
            tf_example = image_example(image_string, label) 
            # and write it into database
            writer.write(tf_example.SerializeToString())

請注意,因爲我們所有的圖像都存儲爲JPEG文件,所以我們使用tf.image.decode_jpeg函數轉換圖像。您還可以將tf.image.decode_image用作通用解碼器。

要檢查所創建數據庫的正確性,可以顯示其中的圖像:

$ python3 show_image.py --path tfrecords/images.tfrecords --method tfrecords

3 加載時間比較

我們將使用來自pexels.com的一些具有不同形狀和jpeg擴展名的開放圖像。並且所有時間測量值將平均進行5000次迭代。此外,平均將減輕OS /硬件特定邏輯(例如數據緩存)的影響。可以預期,正在評估的第一種方法中的第一次迭代將遭受從磁盤到緩存的數據初始加載,而其他方法將沒有這種情況。

所有實驗都針對BGR和RGB圖像模式運行,以涵蓋所有潛在需求和不同任務。請記住,Pillow和Pillow-SIMD不能在同一虛擬環境中使用。爲了創建最終的比較表,我們對Pillow和Pillow-SIMD做了兩個單獨的實驗。

要運行測量,請使用:

python3 benchmark.py --path images/pexels --method cv2 pil turbojpeg lmdb tfrecords --iters 100 --mode BGR

在這裏插入圖片描述
在這裏插入圖片描述
在這裏插入圖片描述
而且,將數據庫的讀取速度與相同的解碼器功能進行比較將很有趣。它可以顯示哪個數據庫更快地加載其數據。在這種情況下,我們對TFRecordsLMDB使用cv2.imdecode函數。

在這裏插入圖片描述
所有實驗的計算依據:

  • Intel® Core™ i7-2600 CPU @ 3.40GHz × 8
  • Ubuntu 16.04 64-bit
  • Python 3.7

4 總結

在此博客中,我們考慮了一些圖像加載方法,並將它們相互比較。JPEG圖像上的比較結果非常有趣。我們可以看到TurboJpeg是將圖像加載爲numpy最快的庫,但是有一個例外——它只能讀取jpeg擴展名的文件。

值得一提的另一件事是,Pillow-SIMD比原始Pillow更快。在我們的任務中,加載速度提高了近40%。

如果您打算使用圖像數據庫,尤其是由於內置的解碼器函數,TFRecords的平均結果要比LMDB更好。另一方面,LMDB使我們可以更快地讀取圖像。當然,您始終可以將解碼器函數和數據庫結合在一起,例如,使用TurboJpeg作爲解碼器,並使用LMDB作爲圖像存儲。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章