cvdatasets/scripts/utils/parser.py

import os
from cvargparse import BaseParser, Arg

from cvdatasets.annotation import AnnotationType

from cvdatasets.utils import read_info_file

DEFAULT_INFO_FILE=os.environ.get("DATA")


def parse_args():

	parser = BaseParser()
	data_help = "YAML file containing dataset and model descriptions"
	dataset_help = "Dataset to display"
	parts_help = "Part annotations to display"

	if DEFAULT_INFO_FILE is None:

		parser.add_args([
			Arg("data", help=data_help),
			Arg("dataset", help=dataset_help),
			Arg("parts", help=parts_help),
		], group_name="Dataset arguments")
	else:
		info_file = read_info_file(DEFAULT_INFO_FILE)
		parser.add_args([
			Arg("data", default=DEFAULT_INFO_FILE,
				help=data_help),
			Arg("dataset", choices=info_file.DATASETS.keys(),
				help=dataset_help),
			Arg("parts", choices=info_file.PART_TYPES.keys(),
				help=parts_help),
		], group_name="Dataset arguments")


	parser.add_args([
		Arg("--subset", "-sub",
			help="Possible subsets: train, test",
			choices=["train", "test"],
			default="train", type=str),

		Arg("--start", "-s",
			help="Image id to start with",
			type=int, default=0),

		Arg("--n_images", "-n",
			help="Number of images to display",
			type=int, default=10),

		Arg("--only_class",
			help="display only the given class",
			type=int, default=-1),
	], group_name="Sample selection arguments")

	parser.add_args([
		Arg("--feature_model", "-fm",
			choices=["inception", "inception_tf", "resnet"]),
		Arg("--features",
			help="pre-extracted train and test features",
			default=[None, None],
			nargs=2, type=str),
	], group_name="Feature arguments")

	parser.add_args([
		Arg("--rnd",
			help="select random subset of present parts",
			action="store_true"),

		Arg("--no_bboxes",
			help="Do not display bounding boxes",
			action="store_true"),

		Arg("--no_parts",
			help="Do not display parts",
			action="store_true"),

		Arg("--crop_to_bb",
			help="Crop image to the bounding box",
			action="store_true"),

		Arg("--no_center_crop", action="store_true"),

		Arg("--crop_uniform",
			help="Try to extend the bounding box to same height and width",
			action="store_true"),

		Arg("--parts_in_bb",
			help="Only display parts, that are inside the bounding box",
			action="store_true"),


		Arg("--ratio",
			help="Part extraction ratio",
			type=float, default=.2),

		Arg("--rescale_size",
			help="rescales the part positions from this size to original image size",
			type=int, default=-1),

		Arg("--uniform_parts", "-u",
			help="Do not use GT parts, but sample parts uniformly from the image",
			action="store_true"),
	], group_name="Display options")

	parser.add_args([
		Arg('--seed', type=int, default=12311123, help='random seed')
	])
	return parser.parse_args()