Skip to content

Commit

Permalink
cleanup download utils
Browse files Browse the repository at this point in the history
  • Loading branch information
pmeier committed Feb 14, 2024
1 parent a52607e commit 1930e30
Showing 1 changed file with 5 additions and 17 deletions.
22 changes: 5 additions & 17 deletions torchvision/datasets/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
import urllib.error
import urllib.request
import zipfile
from typing import Any, Callable, Dict, IO, Iterable, Iterator, List, Optional, Tuple, TypeVar, Union
from typing import Any, Callable, Dict, IO, Iterable, List, Optional, Tuple, TypeVar, Union
from urllib.parse import urlparse

import numpy as np
Expand All @@ -24,24 +24,12 @@
USER_AGENT = "pytorch/vision"


def _save_response_content(
content: Iterator[bytes],
destination: Union[str, pathlib.Path],
length: Optional[int] = None,
) -> None:
with open(destination, "wb") as fh, tqdm(total=length) as pbar:
for chunk in content:
# filter out keep-alive new chunks
if not chunk:
continue

fh.write(chunk)
pbar.update(len(chunk))


def _urlretrieve(url: str, filename: Union[str, pathlib.Path], chunk_size: int = 1024 * 32) -> None:
with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": USER_AGENT})) as response:
_save_response_content(iter(lambda: response.read(chunk_size), b""), filename, length=response.length)
with open(filename, "wb") as fh, tqdm(total=response.length) as pbar:
while chunk := response.read(chunk_size):
fh.write(chunk)
pbar.update(len(chunk))


def calculate_md5(fpath: Union[str, pathlib.Path], chunk_size: int = 1024 * 1024) -> str:
Expand Down

0 comments on commit 1930e30

Please sign in to comment.