Using datasets.utils extract function; accounting for train_extra dataset unzipping (#1068)

torchvision/datasets/cityscapes.py

@@ -3,6 +3,7 @@ import os
 from collections import namedtuple
 import zipfile
 
+from .utils import extract_archive
 from .vision import VisionDataset
 from PIL import Image
 
@@ -126,6 +127,9 @@ class Cityscapes(VisionDataset):
                              ' or "color"')
 
         if not os.path.isdir(self.images_dir) or not os.path.isdir(self.targets_dir):
+            if split == 'train_extra':
+                image_dir_zip = os.path.join(self.root, 'leftImg8bit') + '_trainextra.zip'
+            else:
                 image_dir_zip = os.path.join(self.root, 'leftImg8bit') + '_trainvaltest.zip'
 
             if self.mode == 'gtFine':
@@ -134,8 +138,8 @@ class Cityscapes(VisionDataset):
                 target_dir_zip = os.path.join(self.root, self.mode)
 
             if os.path.isfile(image_dir_zip) and os.path.isfile(target_dir_zip):
-                extract_cityscapes_zip(zip_location=image_dir_zip, root=self.root)
-                extract_cityscapes_zip(zip_location=target_dir_zip, root=self.root)
+                extract_archive(from_path=image_dir_zip, to_path=self.root)
+                extract_archive(from_path=target_dir_zip, to_path=self.root)
             else:
                 raise RuntimeError('Dataset not found or incomplete. Please make sure all required folders for the'
                                    ' specified "split" and "mode" are inside the "root" directory')
@@ -201,9 +205,3 @@ class Cityscapes(VisionDataset):
             return '{}_color.png'.format(mode)
         else:
             return '{}_polygons.json'.format(mode)
-
-
-def extract_cityscapes_zip(zip_location, root):
-    zip_file = zipfile.ZipFile(zip_location, 'r')
-    zip_file.extractall(root)
-    zip_file.close()