Skip to content

Commit

Permalink
add dataset tag and label filters when exporting a dataset from cli
Browse files Browse the repository at this point in the history
  • Loading branch information
Rustem Galiullin committed Aug 15, 2023
1 parent d3761f2 commit 2309cce
Showing 1 changed file with 26 additions and 0 deletions.
26 changes: 26 additions & 0 deletions fiftyone/core/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
|
"""
import argparse
import warnings
from collections import defaultdict
from datetime import datetime
import json
Expand Down Expand Up @@ -815,6 +816,18 @@ def setup(parser):
"`fiftyone.core.collections.SampleCollection.export()`"
),
)
parser.add_argument(
"--filters",
nargs="+",
metavar="KEY=VAL",
action=_ParseKwargsAction,
help=(
"Sample tags or class labels to filter dataset. "
"To use sample tags, pass tags as `tags=train,val` and "
"to use label filters, pass label field and values as in "
"ground_truth=car,person,dog"
),
)

@staticmethod
def execute(parser, args):
Expand All @@ -824,9 +837,22 @@ def execute(parser, args):
dataset_type = etau.get_class(args.type) if args.type else None
label_field = args.label_field
kwargs = args.kwargs or {}
label_filters = args.filters or {}
tags = label_filters.pop("tags", [])

dataset = fod.load_dataset(name)

if tags:
dataset = dataset.match_tags(tags)
for k, v in label_filters.items():
if not dataset.has_sample_field(k):
warnings.warn(
f"Dataset {dataset.name!r} does not contain label field {k!r}"
)
continue
filter_fn = fo.ViewField("label").is_in(v)
dataset = dataset.filter_labels(k, filter_fn)

if json_path:
dataset.write_json(json_path)
print("Dataset '%s' exported to '%s'" % (name, json_path))
Expand Down

0 comments on commit 2309cce

Please sign in to comment.