{"cells":[{"metadata":{},"cell_type":"markdown","source":"# OpenCV+albumentations is faster than PIL+torchvision\nYou can speed up pytorch.utils.data.DataLoader by using cv2.imread + albumentations, instead of PIL.Image.open + torchvision.transforms."},{"metadata":{},"cell_type":"markdown","source":"Based on an official report [here](https://github.com/albu/albumentations#benchmarking-results), this kernel shows some speed comparisons between cv2 method and pil method. Most of cases, cv2 method is faster than pil one."},{"metadata":{},"cell_type":"markdown","source":"The work is now in progress. Feel free to post comments or suggestions. Thanks!"},{"metadata":{},"cell_type":"markdown","source":"# Some Settings"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport gc\nimport time\n\nimport cv2\nimport albumentations as A\nimport numpy as np\nfrom albumentations.pytorch import ToTensor\nfrom PIL import Image\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\n\n\nINPUT_DIR = '../input/all-dogs/all-dogs'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_files = [os.path.join(INPUT_DIR, f) for f in os.listdir(INPUT_DIR)]\nprint('total image files: {}'.format(len(image_files)))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def test(f, n_trials=5):\n    elapsed_times = []\n    for i in range(n_trials):\n        t1 = time.time()\n        f()\n        t2 = time.time()\n        elapsed_times.append(t2-t1)\n    print('Mean: {:.3f}s - Std: {:.3f}s - Max: {:.3f}s - Min: {:.3f}s'.format(\n        np.mean(elapsed_times),\n        np.std(elapsed_times),\n        np.max(elapsed_times),\n        np.min(elapsed_times)\n    ))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Simple Comparisons\nDue to memory constraint, using only 1,000 images."},{"metadata":{"trusted":true},"cell_type":"code","source":"image_files_1000 = image_files[:1000]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Loading Speed\n- cv2: cv2.imread\n- PIL: PIL.Image.open"},{"metadata":{"trusted":true},"cell_type":"code","source":"def cv2_imread(path):\n    img = cv2.imread(path)\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    return img\n\nf = lambda: [cv2_imread(f) for f in image_files_1000]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test(f, n_trials=5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"f = lambda: [Image.open(f) for f in image_files_1000]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test(f, n_trials=5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# use them following test\n\ncv2_img_1000 = [cv2_imread(f) for f in image_files_1000]\npil_img_1000 = [Image.open(f) for f in image_files_1000]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Resize Speed\nresize images to 64x64 by bilinear-interpolation.\n- cv2: albumentations.Resize\n- PIL: torchvision.transforms.Resize"},{"metadata":{"trusted":true},"cell_type":"code","source":"cv2_transform = A.Compose([A.Resize(64, 64, interpolation=cv2.INTER_LINEAR)])\n\nf = lambda: [cv2_transform(image=img)['image'] for img in cv2_img_1000]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test(f, n_trials=5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pil_transform = transforms.Compose([transforms.Resize((64, 64), interpolation=2)])\n\nf = lambda: [pil_transform(img) for img in pil_img_1000]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test(f, n_trials=5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# use them following test\n\ncv2_img_1000 = [cv2_transform(image=img)['image'] for img in cv2_img_1000]\npil_img_1000 = [pil_transform(img) for img in pil_img_1000]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## ToTensor Speed\n- cv2: albumentations.pytorch.ToTensor\n- PIL: torchvision.ToTensor"},{"metadata":{"trusted":true},"cell_type":"code","source":"cv2_transform = A.Compose([ToTensor()])\n\nf = lambda: [cv2_transform(image=img)['image'] for img in cv2_img_1000]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test(f, n_trials=5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pil_transform = transforms.Compose([transforms.ToTensor()])\n\nf = lambda: [pil_transform(img) for img in pil_img_1000]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test(f, n_trials=5)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# DataLoader Comparisons\nDue to time constraints, using only 1,000 images."},{"metadata":{"trusted":true},"cell_type":"code","source":"image_files_1000 = image_files[:1000]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Assumptions\n- Using pytorch. You can get image data as torch.Tensor through torch.utils.data.DataLoader.\n- Limited memory size. You cannot load all image data on memory so use load method.\n- Image size differs. MUST Resize."},{"metadata":{"trusted":true},"cell_type":"code","source":"class BaseDataset(Dataset):\n    def __init__(self, files, transform=None):\n        super().__init__()\n        self.files = files\n        self.transform = transform\n    \n    def __len__(self):\n        return len(self.files)\n\n\nclass PILDataset(BaseDataset):\n    def __getitem__(self, idx):\n        file = self.files[idx]\n        img = Image.open(file)\n        if self.transform is not None:\n            img = self.transform(img)\n            \n        return img\n\n    \nclass CV2Dataset(BaseDataset):\n    def __getitem__(self, idx):\n        file = self.files[idx]\n        img = cv2.imread(file)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        if self.transform is not None:\n            img = self.transform(image=img)['image']\n            \n        return img","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def dataloader_test(files, transform, test_type='cv2', batch_size=64, n_trials=5):\n    assert test_type in ['cv2', 'pil']\n    \n    if test_type == 'cv2':\n        test_dataset = CV2Dataset(files, transform=transform)\n        test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)\n    else:\n        test_dataset = PILDataset(files, transform=transform)\n        test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)\n    \n    def f():\n        for batch in test_dataloader:\n            pass\n    \n    test(f, n_trials=n_trials)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Resize-ToTensor"},{"metadata":{"trusted":true},"cell_type":"code","source":"cv2_transform = A.Compose([\n    A.Resize(64, 64, interpolation=cv2.INTER_LINEAR),\n    ToTensor()\n])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataloader_test(image_files_1000, cv2_transform, test_type='cv2', batch_size=64, n_trials=5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pil_transform = transforms.Compose([\n    transforms.Resize((64, 64), interpolation=2),\n    transforms.ToTensor()\n])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataloader_test(image_files_1000, pil_transform, test_type='pil', batch_size=64, n_trials=5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"## Resize-CenterCrop-ToTensor"},{"metadata":{"trusted":true},"cell_type":"code","source":"cv2_transform = A.Compose([\n    A.SmallestMaxSize(64, interpolation=cv2.INTER_LINEAR),\n    A.CenterCrop(64, 64),\n    ToTensor()\n])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataloader_test(image_files_1000, cv2_transform, test_type='cv2', batch_size=64, n_trials=5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pil_transform = transforms.Compose([\n    transforms.Resize(64, interpolation=2),\n    transforms.CenterCrop(64),\n    transforms.ToTensor()\n])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataloader_test(image_files_1000, pil_transform, test_type='pil', batch_size=64, n_trials=5)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Resize-HorizontalFlip-ToTensor"},{"metadata":{"trusted":true},"cell_type":"code","source":"cv2_transform = A.Compose([\n    A.Resize(64, 64, interpolation=cv2.INTER_LINEAR),\n    A.HorizontalFlip(p=0.5),\n    ToTensor()\n])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataloader_test(image_files_1000, cv2_transform, test_type='cv2', batch_size=64, n_trials=5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pil_transform = transforms.Compose([\n    transforms.Resize((64, 64), interpolation=2),\n    transforms.RandomHorizontalFlip(p=0.5),\n    transforms.ToTensor()\n])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataloader_test(image_files_1000, pil_transform, test_type='pil', batch_size=64, n_trials=5)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Resize-RandomCrop-ToTensor"},{"metadata":{"trusted":true},"cell_type":"code","source":"cv2_transform = A.Compose([\n    A.Resize(96, 96, interpolation=cv2.INTER_LINEAR),\n    A.RandomCrop(64, 64),\n    ToTensor()\n])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataloader_test(image_files_1000, cv2_transform, test_type='cv2', batch_size=64, n_trials=5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pil_transform = transforms.Compose([\n    transforms.Resize((96, 96), interpolation=2),\n    transforms.RandomCrop(64),\n    transforms.ToTensor()\n])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataloader_test(image_files_1000, pil_transform, test_type='pil', batch_size=64, n_trials=5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}