python多线程读图片

用 threading 模块多线程读图片加速，flickr25k 和 nuswide 两个数据集的图片准备见 [1-4]，图像预处理程序来自 [5]。

为了测试，写了一个叫 LazyImage 的类，和其多线程版本 LazyImage_MT，代码：

import time
import multiprocessing
import threading
import os
import os.path as osp
import numpy as np
from PIL import Image

#
# 普通顺序加载：LazyImage、ImageF25k、ImageNUS
#

class LazyImage:
    """mimics np.ndarray, but uses lazy loading"""

    def __init__(self, image_path, image_size=224):
        """image_size: int"""
        self.image_path = image_path
        self.image_size = image_size
        # used in resizing
        self.lower_half = image_size // 2
        self.upper_half = (image_size + 1) // 2

    def __getitem__(self, index):
        if isinstance(index, int):
            return self._load_image(index)
        elif isinstance(index, (np.ndarray, list, tuple)):
            if isinstance(index, np.ndarray):
                assert 1 == index.ndim, "* index should be vector"
            return np.vstack([np.expand_dims(
                self._load_image(i), 0) for i in index])

        raise NotImplemented

    def _load_image(self, full_index):
        """loads single image & resizes
        Input:
            - full_index: int, the sample ID
        """
        img = Image.open(self._get_image_path(full_index))
        xsize, ysize = img.size
        seldim = min(xsize, ysize)
        rate = float(self.image_size) / seldim
        img = img.resize((int(xsize * rate), int(ysize * rate)))
        nxsize, nysize = img.size
        cx, cy = nxsize / 2.0, nysize / 2.0
        box = (cx - self.lower_half, cy - self.lower_half, cx + self.upper_half, cy + self.upper_half)
        img = img.crop(box)
        img = img.convert("RGB")  # can deal with grey-scale images
        img = img.resize((224, 224))
        img = np.array(img, dtype=np.float32)
        return img  # [H, W, C]

    def _get_image_path(self, full_index):
        """get image path according to sample ID"""
        raise NotImplemented


class ImageF25k(LazyImage):
    def _get_image_path(self, full_index):
        # shift to 1-base
        return osp.join(self.image_path, "im{}.jpg".format(full_index + 1))


class ImageNUS(LazyImage):
    """depends on (github) iTomxy/data/nuswide/make.image.link.py"""
    def _get_image_path(self, full_index):
        # remain 0-base as is
        return osp.join(self.image_path, "{}.jpg".format(full_index))

#
# 多线程加载：LazyImage_MT、ImageF25k_MT、ImageNUS_MT
#

class LazyImage_MT:
    """multi-threading version"""

    def __init__(self, image_path, image_size=224, n_thread=None):
        """image_size: int"""
        self.image_path = image_path
        self.image_size = image_size
        # used in resizing
        self.lower_half = image_size // 2
        self.upper_half = (image_size + 1) // 2
        self._mutex_put = threading.Lock()
        self._buffer = []
        self.n_thread = n_thread if (n_thread is not None) else \
            max(1, multiprocessing.cpu_count() - 2)

    def __getitem__(self, index):
        if isinstance(index, int):
            return self._load_image(index)
        elif isinstance(index, (np.ndarray, list, tuple)):
            if isinstance(index, np.ndarray):
                assert 1 == index.ndim, "* index should be vector"
            if self.n_thread < 2:
                return np.vstack([np.expand_dims(self._load_image(i), 0) for i in index])

            self._buffer = []
            batch_size = (len(index) + self.n_thread - 1) // self.n_thread
            t_list = []
            for tid in range(self.n_thread):
                t = threading.Thread(target=self._load_image_mt, args=(
                    index, range(tid * batch_size, min((tid + 1) * batch_size, len(index)))))
                t_list.append(t)
                t.start()
            for t in t_list:
                t.join()
            del t_list
            assert len(self._buffer) == len(index)
            self._buffer = [t[1] for t in sorted(self._buffer, key=lambda _t: _t[0])]
            return np.vstack(self._buffer)

        raise NotImplemented

    def _load_image_mt(self, indices, seg_meta_indices):
        batch_images = [(mid, np.expand_dims(self._load_image(indices[mid]), 0))
            for mid in seg_meta_indices]
        self._mutex_put.acquire()
        self._buffer.extend(batch_images)
        self._mutex_put.release()

    def _load_image(self, full_index):
        """loads single image & resizes
        Input:
            - full_index: int, the sample ID
        """
        img = Image.open(self._get_image_path(full_index))
        xsize, ysize = img.size
        seldim = min(xsize, ysize)
        rate = float(self.image_size) / seldim
        img = img.resize((int(xsize * rate), int(ysize * rate)))
        nxsize, nysize = img.size
        cx, cy = nxsize / 2.0, nysize / 2.0
        box = (cx - self.lower_half, cy - self.lower_half, cx + self.upper_half, cy + self.upper_half)
        img = img.crop(box)
        img = img.convert("RGB")  # can deal with grey-scale images
        img = img.resize((224, 224))
        img = np.array(img, dtype=np.float32)
        return img  # [H, W, C]

    def _get_image_path(self, full_index):
        """get image path according to sample ID"""
        raise NotImplemented


class ImageF25k_MT(LazyImage_MT):
    def _get_image_path(self, full_index):
        # shift to 1-base
        return osp.join(self.image_path, "im{}.jpg".format(full_index + 1))


class ImageNUS_MT(LazyImage_MT):
    """depends on (github) iTomxy/data/nuswide/make.image.link.py"""
    def _get_image_path(self, full_index):
        # remain 0-base as is
        return osp.join(self.image_path, "{}.jpg".format(full_index))

#
# 测试：速度、一致性
#

if "__main__" == __name__:
    batch_size = 128
    N_F25K = 25000
    N_NUS = 269648
    indices_f25k = np.arange(N_F25K)
    indices_nus = np.arange(N_NUS)

    print("-> flickr25k")
    tic = time.time()
    im_f25k = ImageF25k("data/flickr25k/mirflickr")
    for i in range(0, N_F25K, batch_size):
        _ = im_f25k[indices_f25k[i: i + batch_size]]
    print(time.time() - tic)  # 214.9833734035492
    # del im_f25k

    print("-> flickr25k multi-threading")
    tic = time.time()
    im_f25k_mt = ImageF25k_MT("data/flickr25k/mirflickr")
    for i in range(0, N_F25K, batch_size):
        _ = im_f25k_mt[indices_f25k[i: i + batch_size]]
    print(time.time() - tic)  # 56.653871297836304
    # del im_f25k_mt

    print("-> nuswide")
    tic = time.time()
    im_nus = ImageNUS("data/nuswide-tc21/images")
    for i in range(0, N_NUS, batch_size):
        _ = im_nus[indices_nus[i: i + batch_size]]
    print(time.time() - tic)  # 631.8568336963654
    # del im_nus

    print("-> nuswide multi-threading")
    tic = time.time()
    im_nus_mt = ImageNUS_MT("data/nuswide-tc21/images")
    for i in range(0, N_NUS, batch_size):
        _ = im_nus_mt[indices_nus[i: i + batch_size]]
    print(time.time() - tic)  # 207.77122569084167
    # del im_nus_mt

    print("-> consistency")
    for i in range(0, N_F25K, batch_size):
        i_s = im_f25k[indices_f25k[i: i + batch_size]]
        i_mt = im_f25k_mt[indices_f25k[i: i + batch_size]]
        print("f25k diff:", (i_s != i_mt).sum())  # 0
        i_s = im_nus[indices_nus[i: i + batch_size]]
        i_mt = im_nus_mt[indices_nus[i: i + batch_size]]
        print("nus diff:", (i_s != i_mt).sum())  # 0
        break

• 输出

-> flickr25k
214.9833734035492
-> flickr25k multi-threading
56.653871297836304
-> nuswide
631.8568336963654
-> nuswide multi-threading
207.77122569084167
-> consistency
f25k diff: 0
nus diff: 0

References

1. MIR-Flickr25K数据集预处理
2. iTomxy/data/flickr25k
3. NUS-WIDE数据集预处理
4. iTomxy/data/nuswide
5. DeXie0808/GCH/load_data.py