Skip to content

use helper function to extract archive in CelebA #4557

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Oct 7, 2021
Merged

use helper function to extract archive in CelebA #4557

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4a44223
use helper function to extract archive in CelebA
pmeier Oct 7, 2021
af72ce1
Merge branch 'main' into fix-celeba
prabhat00155 Oct 7, 2021
3e646c6
Merge branch 'main' into fix-celeba
prabhat00155 Oct 7, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 2 additions & 5 deletions torchvision/datasets/celeba.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import PIL
import torch

from .utils import download_file_from_google_drive, check_integrity, verify_str_arg
from .utils import download_file_from_google_drive, check_integrity, verify_str_arg, extract_archive
from .vision import VisionDataset

CSV = namedtuple("CSV", ["header", "index", "data"])
Expand Down Expand Up @@ -142,17 +142,14 @@ def _check_integrity(self) -> bool:
return os.path.isdir(os.path.join(self.root, self.base_folder, "img_align_celeba"))

def download(self) -> None:
import zipfile

if self._check_integrity():
print("Files already downloaded and verified")
return

for (file_id, md5, filename) in self.file_list:
download_file_from_google_drive(file_id, os.path.join(self.root, self.base_folder), filename, md5)

with zipfile.ZipFile(os.path.join(self.root, self.base_folder, "img_align_celeba.zip"), "r") as f:
f.extractall(os.path.join(self.root, self.base_folder))
extract_archive(os.path.join(self.root, self.base_folder, "img_align_celeba.zip"))

def __getitem__(self, index: int) -> Tuple[Any, Any]:
X = PIL.Image.open(os.path.join(self.root, self.base_folder, "img_align_celeba", self.filename[index]))
Expand Down