Read COCO dataset from ZIP file

test/test_datasets_utils.py

@@ -3,6 +3,7 @@
 import tempfile
 import torchvision.datasets.utils as utils
 import unittest
+import zipfile
 
 TEST_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                          'assets', 'grace_hopper_517x606.jpg')
@@ -41,6 +42,31 @@ def test_download_url_retry_http(self):
         assert not len(os.listdir(temp_dir)) == 0, 'The downloaded root directory is empty after download.'
         shutil.rmtree(temp_dir)
 
+    def test_ziplookup(self):
+        temp_dir = tempfile.mkdtemp()
+        temp_filename = os.path.join(temp_dir, "ziplookup.zip")
+        try:
+            z = zipfile.ZipFile(temp_filename, "w", zipfile.ZIP_STORED, allowZip64=True)
+            z.write(TEST_FILE, "hopper.jpg")
+            z.write(TEST_FILE)
+            z.write(TEST_FILE, "hopper79.jpg")
+            z.write(TEST_FILE, "somepath/hopper.jpg")
+            z.close()
+
+            lookup = utils.ZipLookup(temp_filename)
+            f = lookup["hopper.jpg"]
+            assert f.name.endswith(".jpg")
+            f = lookup["somepath/hopper.jpg"]
+            assert f.name.endswith(".jpg")
+            try:
+                f = lookup["does_not_exist.jpg"]
+                assert False, "Should not return something for non-existant file"
+            except KeyError:
+                pass
+            del lookup
+        finally:
+            shutil.rmtree(temp_dir)
+
 
 if __name__ == '__main__':
     unittest.main()

torchvision/datasets/celeba.py

@@ -53,7 +53,8 @@ def __init__(self, root,
                  transform=None, target_transform=None,
                  download=False):
         import pandas
-        super(CelebA, self).__init__(root)
+        super(CelebA, self).__init__(root,
+                                     root_zipfilename=os.path.join(self.root, self.base_folder, "img_align_celeba.zip"))
         self.split = split
         if isinstance(target_type, list):
             self.target_type = target_type
@@ -131,7 +132,11 @@ def download(self):
             f.extractall(os.path.join(self.root, self.base_folder))
 
     def __getitem__(self, index):
-        X = PIL.Image.open(os.path.join(self.root, self.base_folder, "img_align_celeba", self.filename[index]))
+        if self.root_zip is not None:
+            f = self.root_zip[os.path.join("img_align_celeba", self.filename[index])]
+            X = PIL.Image.open(f)
+        else:
+            X = PIL.Image.open(os.path.join(self.root, self.base_folder, "img_align_celeba", self.filename[index]))
 
         target = []
         for t in self.target_type:

torchvision/datasets/coco.py

@@ -44,7 +44,8 @@ class CocoCaptions(VisionDataset):
     """
 
     def __init__(self, root, annFile, transform=None, target_transform=None, transforms=None):
-        super(CocoCaptions, self).__init__(root, transforms, transform, target_transform)
+        super(CocoCaptions, self).__init__(root, transforms, transform, target_transform,
+                                           root_zipfilename=self.root + ".zip")
         from pycocotools.coco import COCO
         self.coco = COCO(annFile)
         self.ids = list(sorted(self.coco.imgs.keys()))
@@ -65,7 +66,11 @@ def __getitem__(self, index):
 
         path = coco.loadImgs(img_id)[0]['file_name']
 
-        img = Image.open(os.path.join(self.root, path)).convert('RGB')
+        if self.root_zip is not None:
+            f = self.root_zip[os.path.split(self.root)[1] + "/" + path]
+            img = Image.open(f).convert('RGB')
+        else:
+            img = Image.open(os.path.join(self.root, path)).convert('RGB')
 
         if self.transforms is not None:
             img, target = self.transforms(img, target)
@@ -89,7 +94,8 @@ class CocoDetection(VisionDataset):
     """
 
     def __init__(self, root, annFile, transform=None, target_transform=None, transforms=None):
-        super(CocoDetection, self).__init__(root, transforms, transform, target_transform)
+        super(CocoDetection, self).__init__(root, transforms, transform, target_transform,
+                                            root_zipfilename=self.root + ".zip")
         from pycocotools.coco import COCO
         self.coco = COCO(annFile)
         self.ids = list(sorted(self.coco.imgs.keys()))
@@ -109,7 +115,11 @@ def __getitem__(self, index):
 
         path = coco.loadImgs(img_id)[0]['file_name']
 
-        img = Image.open(os.path.join(self.root, path)).convert('RGB')
+        if self.root_zip is not None:
+            f = self.root_zip[os.path.split(self.root)[1] + "/" + path]
+            img = Image.open(f).convert('RGB')
+        else:
+            img = Image.open(os.path.join(self.root, path)).convert('RGB')
         if self.transforms is not None:
             img, target = self.transforms(img, target)
 

torchvision/datasets/utils.py

@@ -3,6 +3,9 @@
 import hashlib
 import errno
 from torch.utils.model_zoo import tqdm
+import zipfile
+import io
+import struct
 
 
 def gen_bar_updater():
@@ -189,3 +192,37 @@ def _save_response_content(response, destination, chunk_size=32768):
                 progress += len(chunk)
                 pbar.update(progress - pbar.n)
         pbar.close()
+
+
+# thread-safe/multiprocessing-safe (unlike a Python ZipFile instance)
+class ZipLookup(object):
+    def __init__(self, filename):
+        self.root_zip_filename = filename
+        self.root_zip_lookup = {}
+
+        with zipfile.ZipFile(filename, "r") as root_zip:
+            for info in root_zip.infolist():
+                if info.filename[-1] == '/':
+                    # skip directories
+                    continue
+                if info.compress_type != zipfile.ZIP_STORED:
+                    raise ValueError("Only uncompressed ZIP file supported: " + info.filename)
+                if info.compress_size != info.file_size:
+                    raise ValueError("Must be the same when uncompressed")
+                self.root_zip_lookup[info.filename] = (info.header_offset, info.compress_size)
+
+    def __getitem__(self, path):
+        z = open(self.root_zip_filename, "rb")
+        header_offset, size = self.root_zip_lookup[path]
+
+        z.seek(header_offset)
+        fheader = z.read(zipfile.sizeFileHeader)
+        fheader = struct.unpack(zipfile.structFileHeader, fheader)
+        offset = header_offset + zipfile.sizeFileHeader + fheader[zipfile._FH_FILENAME_LENGTH] + \
+            fheader[zipfile._FH_EXTRA_FIELD_LENGTH]
+
+        z.seek(offset)
+        f = io.BytesIO(z.read(size))
+        f.name = path
+        z.close()
+        return f

torchvision/datasets/vision.py

@@ -1,12 +1,13 @@
 import os
 import torch
 import torch.utils.data as data
+from .utils import ZipLookup
 
 
 class VisionDataset(data.Dataset):
     _repr_indent = 4
 
-    def __init__(self, root, transforms=None, transform=None, target_transform=None):
+    def __init__(self, root, transforms=None, transform=None, target_transform=None, root_zipfilename=None):
         if isinstance(root, torch._six.string_classes):
             root = os.path.expanduser(root)
         self.root = root
@@ -25,6 +26,12 @@ def __init__(self, root, transforms=None, transform=None, target_transform=None)
             transforms = StandardTransform(transform, target_transform)
         self.transforms = transforms
 
+        self.root_zip = None
+        self.root_zipfilename = root_zipfilename
+        if self.root_zipfilename is not None and os.path.exists(self.root_zipfilename):
+            self.root_zip = ZipLookup(self.root_zipfilename)
+            print("Using ZIP file for data source:", self.root_zipfilename)
+
     def __getitem__(self, index):
         raise NotImplementedError