Skip to content

Commit

Permalink
Add collection filtering to download commands
Browse files Browse the repository at this point in the history
  • Loading branch information
danlamanna committed Feb 11, 2022
1 parent 90e0d5c commit 34d3b89
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 8 deletions.
8 changes: 6 additions & 2 deletions isic_cli/image.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,10 @@ def _download_image(image: dict, to: Path, progress, task) -> None:
def download_images(
ctx: typer.Context,
search: str = typer.Option(''),
collections: str = typer.Option(
'',
help='Limit the images based on a comma separated string of collection ids (see isic collection list).', # noqa: E501
),
max_images: int = typer.Option(1_000, min=0, help='Use a value of 0 to disable the limit.'),
outdir: Path = typer.Option(Path('images'), file_okay=False, dir_okay=True, writable=True),
):
Expand All @@ -56,12 +60,12 @@ def download_images(
outdir.mkdir(exist_ok=True)
with Progress() as progress:
with get_session(ctx.obj.auth_headers) as session:
num_images = get_num_images(session, search)
num_images = get_num_images(session, search, collections)
if max_images > 0:
num_images = min(num_images, max_images)

task = progress.add_task(f'Downloading images ({num_images} total)', total=num_images)
images = get_images(session, search)
images = get_images(session, search, collections)

if max_images > 0:
images = itertools.islice(images, max_images)
Expand Down
8 changes: 6 additions & 2 deletions isic_cli/metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,10 @@ def validate_metadata(csv_file: Path):
def download(
ctx: typer.Context,
search: str = typer.Option(''),
collections: str = typer.Option(
'',
help='Limit the images based on a comma separated string of collection ids (see isic collection list).', # noqa: E501
),
max_results: int = typer.Option(1_000, min=0, help='Use a value of 0 to disable the limit.'),
):
"""
Expand All @@ -92,11 +96,11 @@ def download(
anatom_site_general:*torso AND image_type:dermoscopic
"""
with get_session(ctx.obj.auth_headers) as session:
num_results = get_num_images(session, search)
num_results = get_num_images(session, search, collections)
if max_results > 0:
num_results = min(num_results, max_results)

images = get_images(session, search)
images = get_images(session, search, collections)

if max_results > 0:
images = itertools.islice(images, max_results)
Expand Down
13 changes: 9 additions & 4 deletions isic_cli/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,8 @@ def get_collections(session: IsicCliSession) -> Iterable[dict]:
next_page = r.json()['next']


def get_images(session: IsicCliSession, search: str) -> Iterable[dict]:
next_page = f'images/search/?query={search}'
def get_images(session: IsicCliSession, search: str, collections: str) -> Iterable[dict]:
next_page = f'images/search/?query={search}' + f'&collections={collections}'

while next_page:
r = session.get(next_page)
Expand All @@ -23,7 +23,12 @@ def get_images(session: IsicCliSession, search: str) -> Iterable[dict]:
next_page = r.json()['next']


def get_num_images(session: IsicCliSession, search: str) -> int:
r = session.get(f'images/search/?query={search}&limit=1')
def get_num_images(session: IsicCliSession, search: str, collections: str) -> int:
params = {
'query': search,
'collections': collections,
'limit': 1,
}
r = session.get('images/search/', params=params)
r.raise_for_status()
return r.json()['count']

0 comments on commit 34d3b89

Please sign in to comment.